perf: optimize circom proof generation

benchmark/msm/README.md

@@ -93,6 +93,10 @@ bazel run -c opt --//:has_openmp --//:has_rtti --//:has_matplotlib //benchmark/m
 
 ### Uniform points
 
+```shell
+bazel run -c opt --//:has_openmp --//:has_rtti --//:has_matplotlib --config cuda //benchmark/msm:msm_benchmark_gpu -- -k 16 -k 17 -k 18 -k 19 -k 20 -k 21 -k 22 -k 23 --test_set non_uniform
+```
+
 #### On RTX-4090
 
 | Exponent | Tachyon  | Tachyon GPU  |
@@ -110,6 +114,10 @@ bazel run -c opt --//:has_openmp --//:has_rtti --//:has_matplotlib //benchmark/m
 
 ### Non-uniform points
 
+```shell
+bazel run -c opt --//:has_openmp --//:has_rtti --//:has_matplotlib --config cuda //benchmark/msm:msm_benchmark_gpu -- -k 16 -k 17 -k 18 -k 19 -k 20 -k 21 -k 22 -k 23
+```
+
 #### On RTX-4090
 
 | Exponent | Tachyon  | Tachyon GPU  |

tachyon/base/BUILD.bazel

@@ -75,6 +75,18 @@ tachyon_cc_library(
     ],
 )
 
+tachyon_cc_library(
+    name = "maybe_owned",
+    hdrs = ["maybe_owned.h"],
+    deps = [":maybe_owned_traits"],
+)
+
+tachyon_cc_library(
+    name = "maybe_owned_traits",
+    hdrs = ["maybe_owned_traits.h"],
+    deps = ["@com_google_absl//absl/types:span"],
+)
+
 tachyon_cc_library(
     name = "no_destructor",
     hdrs = ["no_destructor.h"],

tachyon/base/buffer/read_only_buffer.cc

@@ -2,8 +2,6 @@
 
 #include <string.h>
 
-#include "tachyon/base/numerics/checked_math.h"
-
 namespace tachyon::base {
 
 bool ReadOnlyBuffer::ReadAt(size_t buffer_offset, uint8_t* ptr,

tachyon/base/buffer/read_only_buffer.h

@@ -12,6 +12,7 @@
 #include "tachyon/base/buffer/copyable_forward.h"
 #include "tachyon/base/endian.h"
 #include "tachyon/base/logging.h"
+#include "tachyon/base/numerics/checked_math.h"
 
 namespace tachyon::base {
 namespace internal {
@@ -133,6 +134,22 @@ class TACHYON_EXPORT ReadOnlyBuffer {
     return true;
   }
 
+  template <typename T>
+  [[nodiscard]] bool ReadPtrAt(size_t buffer_offset, T** ptr,
+                               size_t ptr_num) const {
+    size_t size = sizeof(T) * ptr_num;
+    base::CheckedNumeric<size_t> len = buffer_offset;
+    size_t size_needed;
+    if (!(len + size).AssignIfValid(&size_needed)) return false;
+    if (size_needed > buffer_len_) {
+      return false;
+    }
+    const char* buffer = reinterpret_cast<const char*>(buffer_);
+    *ptr = const_cast<T*>(reinterpret_cast<const T*>(&buffer[buffer_offset]));
+    buffer_offset_ = buffer_offset + size;
+    return true;
+  }
+
   [[nodiscard]] bool Read(uint8_t* ptr, size_t size) const {
     return ReadAt(buffer_offset_, ptr, size);
   }
@@ -142,6 +159,11 @@ class TACHYON_EXPORT ReadOnlyBuffer {
     return ReadAt(buffer_offset_, std::forward<T>(value));
   }
 
+  template <typename T>
+  [[nodiscard]] bool ReadPtr(T** ptr, size_t ptr_num) const {
+    return ReadPtrAt(buffer_offset_, ptr, ptr_num);
+  }
+
   [[nodiscard]] bool Read16BE(uint16_t* ptr) const {
     return Read16BEAt(buffer_offset_, ptr);
   }

tachyon/base/maybe_owned.h

@@ -0,0 +1,105 @@
+#ifndef TACHYON_BASE_MAYBE_OWNED_H_
+#define TACHYON_BASE_MAYBE_OWNED_H_
+
+#include <utility>
+
+#include "tachyon/base/maybe_owned_traits.h"
+
+namespace tachyon::base {
+
+template <typename T>
+class MaybeOwned {
+ public:
+  using Ref = typename MaybeOwnedTraits<T>::Ref;
+  using ConstRef = typename MaybeOwnedTraits<T>::ConstRef;
+  using Ptr = typename MaybeOwnedTraits<T>::Ptr;
+  using ConstPtr = typename MaybeOwnedTraits<T>::ConstPtr;
+
+  constexpr MaybeOwned() : ptr_(MaybeOwnedTraits<T>::Default()) {}
+  constexpr MaybeOwned(const MaybeOwned& other)
+      : value_(other.value_), owned_(other.owned_) {
+    UpdatePtr(other);
+  }
+  constexpr MaybeOwned& operator=(const MaybeOwned& other) {
+    value_ = other.value_;
+    owned_ = other.owned_;
+    UpdatePtr(other);
+    return *this;
+  }
+  constexpr MaybeOwned(MaybeOwned&& other)
+      : value_(std::move(other.value_)),
+        owned_(std::exchange(other.owned_, false)) {
+    UpdatePtr(other);
+    MaybeOwnedTraits<T>::ResetPtr(&other.ptr_);
+  }
+  constexpr MaybeOwned& operator=(MaybeOwned&& other) {
+    value_ = std::move(other.value_);
+    owned_ = std::exchange(other.owned_, false);
+    UpdatePtr(other);
+    MaybeOwnedTraits<T>::ResetPtr(&other.ptr_);
+    return *this;
+  }
+  // NOTE(chokobole): |explicit| keyword is explicitly removed for convenience.
+  //
+  // class Storage {
+  //  public:
+  //   Storage(const std::string& data): data_(data) {}
+  //   Storage(std::string&& data): data_(std::move(data)) {}
+  //  private:
+  //   MaybeOwned<std::string> data_;
+  // };
+  // NOLINTNEXTLINE(runtime/explicit)
+  constexpr MaybeOwned(const T& value) : value_(value), owned_(true) {
+    MaybeOwnedTraits<T>::UpdatePtr(&ptr_, value_);
+  }
+  // NOLINTNEXTLINE(runtime/explicit)
+  constexpr MaybeOwned(T&& value) : value_(std::move(value)), owned_(true) {
+    MaybeOwnedTraits<T>::UpdatePtr(&ptr_, value_);
+  }
+  // NOLINTNEXTLINE(runtime/explicit)
+  constexpr MaybeOwned(Ptr ptr) : ptr_(ptr) {}
+  constexpr MaybeOwned& operator=(const T& value) {
+    value_ = value;
+    owned_ = true;
+    MaybeOwnedTraits<T>::UpdatePtr(&ptr_, value_);
+    return *this;
+  }
+  constexpr MaybeOwned& operator=(T&& value) {
+    value_ = std::move(value);
+    owned_ = true;
+    MaybeOwnedTraits<T>::UpdatePtr(&ptr_, value_);
+    return *this;
+  }
+  constexpr MaybeOwned& operator=(Ptr ptr) {
+    MaybeOwnedTraits<T>::Release(value_);
+    ptr_ = ptr;
+    owned_ = false;
+    return *this;
+  }
+
+  Ref operator*() { return MaybeOwnedTraits<T>::ToRef(ptr_); }
+  ConstRef operator*() const { return MaybeOwnedTraits<T>::ToConstRef(ptr_); }
+
+  Ptr operator->() { return ptr(); }
+  ConstPtr operator->() const { return ptr(); }
+
+  Ptr ptr() { return ptr_; }
+  ConstPtr ptr() const { return ptr_; }
+
+ private:
+  void UpdatePtr(const MaybeOwned& other) {
+    if (owned_) {
+      MaybeOwnedTraits<T>::UpdatePtr(&ptr_, value_);
+    } else {
+      ptr_ = other.ptr_;
+    }
+  }
+
+  T value_;
+  Ptr ptr_;
+  bool owned_ = false;
+};
+
+}  // namespace tachyon::base
+
+#endif  // TACHYON_BASE_MAYBE_OWNED_H_

tachyon/base/maybe_owned_traits.h

@@ -0,0 +1,55 @@
+#ifndef TACHYON_BASE_MAYBE_OWNED_TRAITS_H_
+#define TACHYON_BASE_MAYBE_OWNED_TRAITS_H_
+
+#include <utility>
+#include <vector>
+
+#include "absl/types/span.h"
+
+namespace tachyon::base {
+
+template <typename T>
+struct MaybeOwnedTraits {
+  using Ref = T&;
+  using ConstRef = const T&;
+  using Ptr = T*;
+  using ConstPtr = const T*;
+
+  static T* Default() { return nullptr; }
+
+  static void UpdatePtr(T** ptr, T& value) { *ptr = &value; }
+
+  static void ResetPtr(T** ptr) { *ptr = nullptr; }
+
+  static T& ToRef(T* ptr) { return *ptr; }
+
+  static const T& ToConstRef(const T* ptr) { return *ptr; }
+
+  static void Release(T& value) {}
+};
+
+template <typename T>
+struct MaybeOwnedTraits<std::vector<T>> {
+  using Ref = absl::Span<T>;
+  using ConstRef = absl::Span<const T>;
+  using Ptr = absl::Span<T>;
+  using ConstPtr = absl::Span<const T>;
+
+  static absl::Span<T> Default() { return {}; }
+
+  static void UpdatePtr(absl::Span<T>* ptr, std::vector<T>& value) {
+    *ptr = absl::MakeSpan(value);
+  }
+
+  static void ResetPtr(absl::Span<T>* ptr) { *ptr = absl::Span<T>(); }
+
+  static absl::Span<T> ToRef(absl::Span<T> ptr) { return ptr; }
+
+  static absl::Span<const T> ToConstRef(absl::Span<const T> ptr) { return ptr; }
+
+  static void Release(std::vector<T>& value) { value.clear(); }
+};
+
+}  // namespace tachyon::base
+
+#endif  // TACHYON_BASE_MAYBE_OWNED_TRAITS_H_

tachyon/math/base/BUILD.bazel

@@ -35,7 +35,6 @@ tachyon_cc_library(
         ":arithmetics",
         ":bit_traits_forward",
         "//tachyon/base:bit_cast",
-        "//tachyon/base:compiler_specific",
         "//tachyon/base:endian_utils",
         "//tachyon/base:random",
         "//tachyon/base/buffer:copyable",

tachyon/math/base/big_int.h

@@ -20,7 +20,6 @@
 
 #include "tachyon/base/bit_cast.h"
 #include "tachyon/base/buffer/copyable.h"
-#include "tachyon/base/compiler_specific.h"
 #include "tachyon/base/endian_utils.h"
 #include "tachyon/base/json/json.h"
 #include "tachyon/base/logging.h"
@@ -43,16 +42,12 @@ TACHYON_EXPORT std::string LimbsToString(const uint64_t* limbs,
 TACHYON_EXPORT std::string LimbsToHexString(const uint64_t* limbs,
                                             size_t limb_nums, bool pad_zero);
 
-constexpr size_t LimbsAlignment(size_t x) {
-  return x % 4 == 0 ? 32 : (x % 2 == 0 ? 16 : 8);
-}
-
 }  // namespace internal
 
 // BigInt is a fixed size array of uint64_t, capable of holding up to |N| limbs,
 // designed to support a wide range of big integer arithmetic operations.
 template <size_t N>
-struct ALIGNAS(internal::LimbsAlignment(N)) BigInt {
+struct BigInt {
   uint64_t limbs[N] = {
       0,
   };

tachyon/math/elliptic_curves/msm/algorithms/icicle/BUILD.bazel

@@ -27,6 +27,7 @@ tachyon_cc_library(
     deps = [
         ":icicle_msm_bn254_g1",
         ":icicle_msm_bn254_g2",
+        "//tachyon:export",
         "//tachyon/base:bit_cast",
         "//tachyon/device/gpu:gpu_device_functions",
         "//tachyon/device/gpu:gpu_enums",

tachyon/math/elliptic_curves/msm/algorithms/icicle/icicle_msm.h

@@ -8,6 +8,7 @@
 #include "tachyon/base/bit_cast.h"
 #include "tachyon/device/gpu/gpu_device_functions.h"
 #include "tachyon/device/gpu/gpu_enums.h"
+#include "tachyon/export.h"
 #include "tachyon/math/elliptic_curves/bn/bn254/g1.h"
 #include "tachyon/math/elliptic_curves/bn/bn254/g2.h"
 #include "tachyon/math/elliptic_curves/msm/algorithms/icicle/icicle_msm_bn254_g1.h"
@@ -16,31 +17,49 @@
 
 namespace tachyon::math {
 
+struct TACHYON_EXPORT IcicleMSMOptions {
+  int points_size = 0;
+  int precompute_factor = 1;
+  int c = 0;
+  int bitsize = 0;
+  int large_bucket_factor = 10;
+  int batch_size = 1;
+  bool are_scalars_on_device = false;
+  bool are_scalars_montgomery_form = true;
+  bool are_points_on_device = false;
+  bool are_points_montgomery_form = true;
+  bool are_results_on_device = false;
+  bool is_big_triangle = false;
+  bool is_async = false;
+};
+
 template <typename Point>
 class IcicleMSM {
  public:
   using Curve = typename Point::Curve;
 
-  IcicleMSM(gpuMemPool_t mem_pool, gpuStream_t stream)
+  IcicleMSM(gpuMemPool_t mem_pool, gpuStream_t stream,
+            const IcicleMSMOptions& options = IcicleMSMOptions())
       : mem_pool_(mem_pool), stream_(stream) {
     ::device_context::DeviceContext ctx{stream_, /*device_id=*/0, mem_pool_};
     config_.reset(new ::msm::MSMConfig{
         ctx,
-        /*points_size=*/0,
-        /*precompute_factor=*/1,
-        /*c=*/0,
-        /*bitsize=*/0,
-        /*large_bucket_factor=*/10,
-        /*batch_size=*/1,
-        /*are_scalars_on_device=*/false,
-        /*are_scalars_montgomery_form=*/true,
+        options.points_size,
+        options.precompute_factor,
+        options.c,
+        options.bitsize,
+        options.large_bucket_factor,
+        options.batch_size,
+        options.are_scalars_on_device,
+        options.are_scalars_montgomery_form,
         // TODO(chokobole): Considering KZG commitment, bases can be loaded to
         // the device just once initially.
-        /*are_points_on_device=*/false,
-        /*are_points_montgomery_form=*/true,
-        /*are_results_on_device=*/false,
-        /*is_big_triangle=*/false,
-        /*is_async=*/false});
+        options.are_points_on_device,
+        options.are_points_montgomery_form,
+        options.are_results_on_device,
+        options.is_big_triangle,
+        options.is_async,
+    });
     VLOG(1) << "IcicleMSM is created";
   }
   IcicleMSM(const IcicleMSM& other) = delete;

tachyon/math/elliptic_curves/pairing/BUILD.bazel

@@ -11,7 +11,10 @@ tachyon_cc_library(
 tachyon_cc_library(
     name = "g2_prepared_base",
     hdrs = ["g2_prepared_base.h"],
-    deps = [":ell_coeff"],
+    deps = [
+        ":ell_coeff",
+        "//tachyon/base/strings:string_util",
+    ],
 )
 
 tachyon_cc_library(

tachyon/math/elliptic_curves/pairing/ell_coeff.h

@@ -31,6 +31,12 @@ class EllCoeff {
                             c1_.ToString(), c2_.ToString());
   }
 
+  std::string ToHexString(bool pad_zero = false) const {
+    return absl::Substitute(
+        "{c0: $0, c1: $1, c2: $2}", c0_.ToHexString(pad_zero),
+        c1_.ToHexString(pad_zero), c2_.ToHexString(pad_zero));
+  }
+
  private:
   F c0_;
   F c1_;