boojum-cuda 0.155.2

Boojum-CUDA is a library implementing GPU-accelerated cryptographic functionality for the zkSync prover
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40

#include "common.cuh"
#include "context.cuh"
#include "goldilocks.cuh"
#include "goldilocks_extension.cuh"
#include "memory.cuh"

namespace goldilocks {

using namespace memory;

template <typename T, int INV_BATCH, bool batch_is_full>
DEVICE_FORCEINLINE void batch_inv_registers(const T *inputs, T *fwd_scan_and_outputs, int runtime_batch_size) {
  // If count < grid size, the kernel is inefficient no matter what (because each thread processes just one element)
  // but we should still bail out if a thread has no assigned elems at all.
  T running_prod = T::one();
  int i = 0;
#pragma unroll
  for (; i < INV_BATCH; i++)
    if (batch_is_full || i < runtime_batch_size) {
      fwd_scan_and_outputs[i] = running_prod;
      running_prod = T::mul(running_prod, inputs[i]);
    }

  T inv = T::inv(running_prod);

  i--;
#pragma unroll
  for (; i >= 0; i--) {
    if (batch_is_full || i < runtime_batch_size) {
      const auto input = inputs[i];
      // Isolates and stores this input's inv
      fwd_scan_and_outputs[i] = T::mul(fwd_scan_and_outputs[i], inv);
      // Removes this input's inv contribution
      if (i > 0)
        inv = T::mul(inv, input);
    }
  }
}

} // namespace goldilocks