feat: use multiple GPUs for Pippenger's partition algorithm (PROOF-831)

sxt/multiexp/pippenger2/BUILD

@@ -146,15 +146,27 @@ sxt_cc_component(
         "//sxt/ristretto/random:element",
     ],
     deps = [
+        ":combination",
         ":partition_product",
         ":partition_table_accessor",
         ":reduce",
         "//sxt/base/container:span",
+        "//sxt/base/container:span_utility",
         "//sxt/base/curve:element",
-        "//sxt/execution/async:future",
+        "//sxt/base/device:memory_utility",
+        "//sxt/base/device:property",
+        "//sxt/base/device:state",
+        "//sxt/base/device:stream",
+        "//sxt/base/device:synchronization",
+        "//sxt/base/iterator:index_range_iterator",
+        "//sxt/base/iterator:index_range_utility",
+        "//sxt/base/log",
+        "//sxt/execution/async:coroutine",
+        "//sxt/execution/device:for_each",
         "//sxt/memory/management:managed_array",
         "//sxt/memory/resource:async_device_resource",
         "//sxt/memory/resource:device_resource",
+        "//sxt/memory/resource:pinned_resource",
     ],
 )
 

sxt/multiexp/pippenger2/multiexponentiation.h

@@ -16,32 +16,50 @@
  */
 #pragma once
 
+#include <iterator>
+
 #include "sxt/base/container/span.h"
+#include "sxt/base/container/span_utility.h"
 #include "sxt/base/curve/element.h"
+#include "sxt/base/device/memory_utility.h"
+#include "sxt/base/device/property.h"
+#include "sxt/base/device/state.h"
+#include "sxt/base/device/stream.h"
 #include "sxt/base/error/assert.h"
-#include "sxt/execution/async/future.h"
+#include "sxt/base/iterator/index_range_iterator.h"
+#include "sxt/base/iterator/index_range_utility.h"
+#include "sxt/base/log/log.h"
+#include "sxt/execution/async/coroutine.h"
+#include "sxt/execution/device/for_each.h"
+#include "sxt/execution/device/synchronization.h"
 #include "sxt/memory/management/managed_array.h"
 #include "sxt/memory/resource/async_device_resource.h"
 #include "sxt/memory/resource/device_resource.h"
+#include "sxt/memory/resource/pinned_resource.h"
+#include "sxt/multiexp/pippenger2/combination.h"
 #include "sxt/multiexp/pippenger2/partition_product.h"
 #include "sxt/multiexp/pippenger2/partition_table_accessor.h"
 #include "sxt/multiexp/pippenger2/reduce.h"
 
 namespace sxt::mtxpp2 {
 //--------------------------------------------------------------------------------------------------
-// multiexponentiate
+// multiexponentiate_options
+//--------------------------------------------------------------------------------------------------
+struct multiexponentiate_options {
+  unsigned split_factor = 1;
+  unsigned min_chunk_size = 64;
+  unsigned max_chunk_size = 1024;
+};
+
+//--------------------------------------------------------------------------------------------------
+// multiexponentiate_no_chunks
 //--------------------------------------------------------------------------------------------------
-/**
- * Compute a multi-exponentiation using an accessor to precompute sums of partition groups.
- *
- * This implements the partition part of Pipenger's algorithm. See Algorithm 7 of
- * https://cacr.uwaterloo.ca/techreports/2010/cacr2010-26.pdf
- */
 template <bascrv::element T>
-xena::future<> multiexponentiate(basct::span<T> res, const partition_table_accessor<T>& accessor,
-                                 unsigned element_num_bytes,
-                                 basct::cspan<uint8_t> scalars) noexcept {
+xena::future<>
+multiexponentiate_no_chunks(basct::span<T> res, const partition_table_accessor<T>& accessor,
+                            unsigned element_num_bytes, basct::cspan<uint8_t> scalars) noexcept {
   auto num_outputs = res.size();
+  auto n = scalars.size() / (num_outputs * element_num_bytes);
   auto num_products = num_outputs * element_num_bytes * 8u;
   SXT_DEBUG_ASSERT(
       // clang-format off
@@ -50,10 +68,12 @@ xena::future<> multiexponentiate(basct::span<T> res, const partition_table_acces
   );
 
   // compute bitwise products
+  basl::info("computing {} bitwise multiexponentiation products of length {}", num_products, n);
   memmg::managed_array<T> products(num_products, memr::get_device_resource());
   co_await partition_product<T>(products, accessor, scalars, 0);
 
   // reduce products
+  basl::info("reducing {} products to {} outputs", num_products, num_products);
   basdv::stream stream;
   memr::async_device_resource resource{stream};
   memmg::managed_array<T> res_dev{num_outputs, &resource};
@@ -64,4 +84,116 @@ xena::future<> multiexponentiate(basct::span<T> res, const partition_table_acces
   basdv::async_copy_device_to_host(res, res_dev, stream);
   co_await xendv::await_stream(stream);
 }
+
+//--------------------------------------------------------------------------------------------------
+// complete_multiexponentiation
+//--------------------------------------------------------------------------------------------------
+template <bascrv::element T>
+xena::future<> complete_multiexponentiation(basct::span<T> res, unsigned element_num_bytes,
+                                            basct::cspan<T> partial_products, unsigned num_products,
+                                            unsigned offset) noexcept {
+  auto num_outputs_slice = res.size();
+  auto num_products_slice = num_outputs_slice * element_num_bytes * 8u;
+
+  // combine the partial results
+  memmg::managed_array<T> products_slice{num_products_slice, memr::get_device_resource()};
+  co_await combine_partial<T>(products_slice, partial_products, num_products, offset);
+
+  // reduce the products
+  basdv::stream stream;
+  memr::async_device_resource resource{stream};
+  memmg::managed_array<T> res_dev{num_outputs_slice, &resource};
+  reduce_products<T>(res_dev, stream, products_slice);
+  products_slice.reset();
+
+  // copy result
+  basdv::async_copy_device_to_host(res, res_dev, stream);
+  co_await xendv::await_stream(stream);
+}
+
+//--------------------------------------------------------------------------------------------------
+// multiexponentiate_impl
+//--------------------------------------------------------------------------------------------------
+template <bascrv::element T>
+xena::future<> multiexponentiate_impl(basct::span<T> res,
+                                      const partition_table_accessor<T>& accessor,
+                                      unsigned element_num_bytes, basct::cspan<uint8_t> scalars,
+                                      const multiexponentiate_options& options) noexcept {
+  auto num_outputs = res.size();
+  auto n = scalars.size() / (num_outputs * element_num_bytes);
+  auto num_products = num_outputs * element_num_bytes * 8u;
+  SXT_DEBUG_ASSERT(
+      // clang-format off
+      scalars.size() % (num_outputs * element_num_bytes) == 0
+      // clang-format on
+  );
+
+  // compute bitwise products
+  //
+  // We split the work by groups of generators so that a single chunk will process
+  // all the outputs for those generators. This minimizes the amount of host->device
+  // copying we need to do for the table of precomputed sums.
+  auto [chunk_first, chunk_last] = basit::split(basit::index_range{0, n}
+                                                    .chunk_multiple(16)
+                                                    .min_chunk_size(options.min_chunk_size)
+                                                    .max_chunk_size(options.max_chunk_size),
+                                                options.split_factor);
+  auto num_chunks = std::distance(chunk_first, chunk_last);
+  if (num_chunks == 1) {
+    multiexponentiate_no_chunks(res, accessor, element_num_bytes, scalars);
+    co_return;
+  }
+
+  memmg::managed_array<T> products{num_products * num_chunks, memr::get_pinned_resource()};
+  size_t chunk_index = 0;
+  basl::info("computing {} bitwise multiexponentiation products of length {} using {} chunks",
+             num_products, n, num_chunks);
+  co_await xendv::concurrent_for_each(
+      chunk_first, chunk_last, [&](const basit::index_range& rng) noexcept -> xena::future<> {
+        basl::info("computing {} multiproducts for generators [{}, {}] on device {}", num_products,
+                   rng.a(), rng.b(), basdv::get_device());
+        memmg::managed_array<T> products_dev{num_products, memr::get_device_resource()};
+        auto scalars_slice = scalars.subspan(num_outputs * element_num_bytes * rng.a(),
+                                             rng.size() * num_outputs * element_num_bytes);
+        co_await partition_product<T>(products_dev, accessor, scalars_slice, rng.a());
+        basdv::stream stream;
+        basdv::async_copy_device_to_host(
+            basct::subspan(products, num_products * chunk_index, num_products), products_dev,
+            stream);
+        ++chunk_index;
+        co_await xendv::await_stream(stream);
+      });
+
+  // complete the multi-exponentiation by splitting the remaining work by output
+  auto [output_first, output_last] =
+      basit::split(basit::index_range{0, num_outputs}, options.split_factor);
+  basl::info("reducing products for {} outputs using {} chunks", num_outputs,
+             std::distance(output_first, output_last));
+  co_await xendv::concurrent_for_each(
+      output_first, output_last, [&](const basit::index_range& rng) noexcept -> xena::future<> {
+        basl::info("reducing products for outputs [{}, {}] on device {}", rng.a(), rng.b(),
+                   basdv::get_device());
+        co_await complete_multiexponentiation<T>(res.subspan(rng.a(), rng.size()),
+                                                 element_num_bytes, products, num_products,
+                                                 rng.a() * element_num_bytes * 8u);
+      });
+}
+
+//--------------------------------------------------------------------------------------------------
+// multiexponentiate
+//--------------------------------------------------------------------------------------------------
+/**
+ * Compute a multi-exponentiation using an accessor to precompute sums of partition groups.
+ *
+ * This implements the partition part of Pipenger's algorithm. See Algorithm 7 of
+ * https://cacr.uwaterloo.ca/techreports/2010/cacr2010-26.pdf
+ */
+template <bascrv::element T>
+xena::future<> multiexponentiate(basct::span<T> res, const partition_table_accessor<T>& accessor,
+                                 unsigned element_num_bytes,
+                                 basct::cspan<uint8_t> scalars) noexcept {
+  multiexponentiate_options options;
+  options.split_factor = static_cast<unsigned>(basdv::get_num_devices());
+  return multiexponentiate_impl(res, accessor, element_num_bytes, scalars, options);
+}
 } // namespace sxt::mtxpp2

sxt/multiexp/pippenger2/multiexponentiation.t.cc

@@ -111,6 +111,38 @@ TEST_CASE("we can compute multiexponentiations using a precomputed table of part
     REQUIRE(res[0] == generators[0].value);
     REQUIRE(res[1] == 2u * generators[0].value);
   }
+
+  SECTION("we can split a multi-exponentiation") {
+    multiexponentiate_options options{
+        .split_factor = 2,
+        .min_chunk_size = 16u,
+    };
+    scalars.resize(32);
+    scalars[0] = 1;
+    scalars[16] = 1;
+    auto fut = multiexponentiate_impl<E>(res, *accessor, 1, scalars, options);
+    xens::get_scheduler().run();
+    REQUIRE(fut.ready());
+    REQUIRE(res[0] == generators[0].value + generators[16].value);
+  }
+
+  SECTION("we can split a multi-exponentiation with more than one output") {
+    multiexponentiate_options options{
+        .split_factor = 2,
+        .min_chunk_size = 16u,
+    };
+    scalars.resize(64);
+    scalars[0] = 1;
+    scalars[1] = 2;
+    scalars[32] = 3;
+    scalars[33] = 4;
+    res.resize(2);
+    auto fut = multiexponentiate_impl<E>(res, *accessor, 1, scalars, options);
+    xens::get_scheduler().run();
+    REQUIRE(fut.ready());
+    REQUIRE(res[0] == generators[0].value + 3 * generators[16].value);
+    REQUIRE(res[1] == 2 * generators[0].value + 4 * generators[16].value);
+  }
 }
 
 TEST_CASE("we can compute multiexponentiations with curve-21") {