[FEATURE] Reduction Kernel with Multi Blocks (p.s. I am too lazy to implement thread coarsening...)

Author: root

configs/lib-tests.yml

@@ -9,4 +9,8 @@ OpTest:
     - nInputs: 2048
       divider: 4
   MulReduction:
-    - nInputs: 32
+    - nInputs: 32
+    # # [NOTE] The following calculation results would be wrong?
+    # - nInputs: 64
+    # - nInputs: 128
+    # - nInputs: 512

csrc/include/pmpp/utils/common.cuh

@@ -16,11 +16,11 @@
         do {                                                                  \
             cudaError_t err_ = (err);                                         \
             if (err_ != cudaSuccess) {                                        \
-                fprintf(stderr, "CUDA error at %s:%d code=%d(%s) \"%s\"",     \
-                        __FILE__, __LINE__, err, cudaGetErrorString(err_),    \
-                        #err);                                                \
-                cudaDeviceReset();                                            \
-                throw std::runtime_error("CUDA error");                       \
+                ::fprintf(stderr, "CUDA error at %s:%d code=%d(%s) \"%s\"",   \
+                          __FILE__, __LINE__, err,                            \
+                          ::cudaGetErrorString(err_), #err);                  \
+                ::cudaDeviceReset();                                          \
+                throw ::std::runtime_error("CUDA error");                     \
             }                                                                 \
         } while (0)
 #endif
@@ -52,4 +52,5 @@ __host__ __device__ void initMemory(T* ptr, size_t n, const T& val)
         ptr[i] = val;
     }
 }
+
 }  // namespace pmpp::cuda

csrc/include/pmpp/utils/math.hpp

@@ -15,7 +15,7 @@ namespace pmpp
  */
 template <typename T1, typename T2>
     requires std::is_integral_v<T1> && std::is_integral_v<T2>
-constexpr auto ceilDiv(T1 a, T2 b) -> T1
+[[nodiscard]] constexpr auto ceilDiv(T1 a, T2 b) -> T1
 {
     return T1((a + b - 1) / b);
 }

csrc/lib/ops/reduction/op.cuh

@@ -11,36 +11,43 @@ template <typename ScalarT, typename PredT>
 __global__ void reductionKernel(const ScalarT* in, uint32_t n, ScalarT* out,
                                 const PredT& pred)
 {
-    // Thread index in the block
-    uint32_t bTid = threadIdx.x;
+    uint32_t stride = blockDim.x;
+    uint32_t segmentId = blockIdx.x;
+    uint32_t segmentSize = 2 * stride;
+    // Block thread index
+    uint32_t bTidx = threadIdx.x;
+    // Global thread index
+    uint32_t gTidx = segmentId * segmentSize + bTidx;
+
     extern __shared__ ScalarT shmem[];
 
-    uint32_t stride = blockDim.x;
-    shmem[bTid] = pred(in[bTid], in[bTid + stride]);
+    shmem[bTidx] = pred(in[gTidx], in[gTidx + stride]);
     stride /= 2;
 
     for (; stride >= 1; stride /= 2) {
         __syncthreads();
-        if (bTid < stride) {
-            shmem[bTid] = pred(shmem[bTid], shmem[bTid + stride]);
+        if (bTidx < stride) {
+            shmem[bTidx] = pred(shmem[bTidx], shmem[bTidx + stride]);
         }
     }
-    if (bTid == 0) {
-        out[0] = shmem[0];
+    if (bTidx == 0) {
+        atomicAdd(out, shmem[0]);
     }
 }
 
 template <typename ScalarT, typename PredT>
 [[nodiscard]] auto launchReduction(const ScalarT* in, uint32_t n,
                                    const PredT& pred) -> ScalarT
 {
+    constexpr uint32_t MAX_BLOCK_THREADS = 1024;
+
     ScalarT* d_out;
     cudaMalloc(&d_out, 1 * sizeof(ScalarT));
 
-    uint32_t nTreads = n / 2;
-    dim3 blockDim = {nTreads, 1, 1};
-    dim3 gridDim = {1, 1, 1};
-    uint32_t shmemSize = blockDim.x * sizeof(ScalarT);
+    uint32_t stride = std::min(n / 2, MAX_BLOCK_THREADS);
+    dim3 blockDim = {stride, 1, 1};
+    dim3 gridDim = {ceilDiv(n, stride * 2), 1, 1};
+    uint32_t shmemSize = stride * sizeof(ScalarT);
 
     reductionKernel<<<gridDim, blockDim, shmemSize>>>(in, n, d_out, pred);