[mlir][sparse][gpu] fix F32 bug for SpMV and SpMM

[aartbik]

Summary

The alpha/beta variables, residing on the host, should have the
32-bit or 64-bit width of the result type. It was formerly always
passed as double.

[Peiming]

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

251–260

This is indeed magic :-)

I would suggest change it to make it a little bit safer from name collision

[aartbik]

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

251–260

You mean adding the variable names as parameter too?
Sure, that seems a good idea, perhaps I can even use the append feature for the derived names then....

[Peiming]

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

251–260

Oh I mean adding extra braces around alphaf, etc

But adding the variable name as parameter might be a good idea too.

[aartbik]

even more magic!

[aartbik]

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

251–260

Oh, braces too. I think in that case, the locals go out of scope at the closing }
making the pointers invalid, right?

But see if you like the new magic better. It is really just to avoid repeating the same block of code over and over again, while keeping it local scope

[Peiming]

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

251–260

Yeah, that's why I still keep alphap out of the scope (only variable ending with -f and -d are in local scope).

[Peiming]

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

251–260

NVM, you are right ;-)

[aartbik]

lint

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

static inline cusparseIndexType_t idxTp(int32_t width) {
   switch (width) {
   case 32:
     return CUSPARSE_INDEX_32I;
   default:
     return CUSPARSE_INDEX_64I;
   }
 }
 
+// Some macro magic to get float/double alpha and beta on host.
+#define ALPHABETA(w, alpha, beta)                                              \
+  float (alpha##f) = 1.0, (beta##f) = 1.0;                                     \
+  double (alpha##d) = 1.0, (beta##d) = 1.0;                                    \
+  void *(alpha##p), *(beta##p);                                                \
+  if ((w) == 32) {                                                             \
+    (alpha##p) = reinterpret_cast<void *>(&(alpha##f));                        \
+    (beta##p) = reinterpret_cast<void *>(&(beta##f));                          \
+  } else {                                                                     \
+    (alpha##p) = reinterpret_cast<void *>(&(alpha##d));                        \
+    (beta##p) = reinterpret_cast<void *>(&(beta##d));                          \
+  }
+
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void *
 mgpuCreateSparseEnv(CUstream /*stream*/) {
   cusparseHandle_t handle = nullptr;
   CUSPARSE_REPORT_IF_ERROR(cusparseCreate(&handle))
   return reinterpret_cast<void *>(handle);
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void

 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t mgpuSpMVBufferSize(
     void *h, void *a, void *x, void *y, int32_t dw, CUstream /*stream*/) {
   cusparseHandle_t handle = reinterpret_cast<cusparseHandle_t>(h);
   cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a);
   cusparseDnVecDescr_t vecX = reinterpret_cast<cusparseDnVecDescr_t>(x);
   cusparseDnVecDescr_t vecY = reinterpret_cast<cusparseDnVecDescr_t>(y);
   cudaDataType_t dtp = dataTp(dw);
-  double alpha = 1.0;
-  double beta = 1.0;
+  ALPHABETA(dw, alpha, beta)
   size_t bufferSize = 0;
   CUSPARSE_REPORT_IF_ERROR(cusparseSpMV_bufferSize(
-      handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, matA, vecX, &beta, vecY,
+      handle, CUSPARSE_OPERATION_NON_TRANSPOSE, alphap, matA, vecX, betap, vecY,
       dtp, CUSPARSE_SPMV_ALG_DEFAULT, &bufferSize))
   return bufferSize == 0 ? 1 : bufferSize; // avoid zero-alloc
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSpMV(void *h, void *a, void *x,
                                                    void *y, int32_t dw,
                                                    void *buf,
                                                    CUstream /*stream*/) {
   cusparseHandle_t handle = reinterpret_cast<cusparseHandle_t>(h);
   cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a);
   cusparseDnVecDescr_t vecX = reinterpret_cast<cusparseDnVecDescr_t>(x);
   cusparseDnVecDescr_t vecY = reinterpret_cast<cusparseDnVecDescr_t>(y);
   cudaDataType_t dtp = dataTp(dw);
-  double alpha = 1.0;
-  double beta = 1.0;
+  ALPHABETA(dw, alpha, beta)
   CUSPARSE_REPORT_IF_ERROR(
-      cusparseSpMV(handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, matA, vecX,
-                   &beta, vecY, dtp, CUSPARSE_SPMV_ALG_DEFAULT, buf))
+      cusparseSpMV(handle, CUSPARSE_OPERATION_NON_TRANSPOSE, alphap, matA, vecX,
+                   betap, vecY, dtp, CUSPARSE_SPMV_ALG_DEFAULT, buf))
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t mgpuSpMMBufferSize(
     void *h, void *a, void *b, void *c, int32_t dw, CUstream /*stream*/) {
   cusparseHandle_t handle = reinterpret_cast<cusparseHandle_t>(h);
   cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a);
   cusparseDnMatDescr_t matB = reinterpret_cast<cusparseDnMatDescr_t>(b);
   cusparseDnMatDescr_t matC = reinterpret_cast<cusparseDnMatDescr_t>(c);
   cudaDataType_t dtp = dataTp(dw);
-  double alpha = 1.0;
-  double beta = 1.0;
+  ALPHABETA(dw, alpha, beta)
   size_t bufferSize = 0;
   CUSPARSE_REPORT_IF_ERROR(cusparseSpMM_bufferSize(
       handle, CUSPARSE_OPERATION_NON_TRANSPOSE,
-      CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, matA, matB, &beta, matC, dtp,
+      CUSPARSE_OPERATION_NON_TRANSPOSE, alphap, matA, matB, betap, matC, dtp,
       CUSPARSE_SPMM_ALG_DEFAULT, &bufferSize))
   return bufferSize == 0 ? 1 : bufferSize; // avoid zero-alloc
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSpMM(void *h, void *a, void *b,
                                                    void *c, int32_t dw,
                                                    void *buf,
                                                    CUstream /*stream*/) {
   cusparseHandle_t handle = reinterpret_cast<cusparseHandle_t>(h);
   cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a);
   cusparseDnMatDescr_t matB = reinterpret_cast<cusparseDnMatDescr_t>(b);
   cusparseDnMatDescr_t matC = reinterpret_cast<cusparseDnMatDescr_t>(c);
   cudaDataType_t dtp = dataTp(dw);
-  double alpha = 1.0;
-  double beta = 1.0;
+  ALPHABETA(dw, alpha, beta)
   CUSPARSE_REPORT_IF_ERROR(
       cusparseSpMM(handle, CUSPARSE_OPERATION_NON_TRANSPOSE,
-                   CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, matA, matB, &beta,
+                   CUSPARSE_OPERATION_NON_TRANSPOSE, alphap, matA, matB, betap,
                    matC, dtp, CUSPARSE_SPMM_ALG_DEFAULT, buf))
 }