upd (#2336)

src/array/cuda/sddmm.cuh

@@ -106,8 +106,11 @@ __global__ void SDDMMCooTreeReduceKernel(
       const Idx lhs_add = UseBcast ? __ldg(lhs_off + i) : i;
       const Idx rhs_add = UseBcast ? __ldg(rhs_off + i) : i;
       DType val = 0.;
-      for (int j = tx; j < reduce_size; j += 32)
+      for (int j = tx; j < reduce_size; j += 64) {
         val += lhsoff[lhs_add * reduce_size + j] * rhsoff[rhs_add * reduce_size + j];
+        if (j + 32 < reduce_size)
+          val += lhsoff[lhs_add * reduce_size + j + 32] * rhsoff[rhs_add * reduce_size + j + 32];
+      }
 #pragma unroll
       for (int offset = 16; offset > 0; offset /= 2)
         val += __shfl_down_sync(full_mask, val, offset);