[mlir][sparse][gpu] add result type to spmv, spmm and sddmm gpu libgen path

[K-Wu]

rm redundant namespaces

[K-Wu]

add computeType in tblgen

[K-Wu]

working implementation though a bit ugly

[K-Wu]

rebase origin/main and add sddmm

[K-Wu]

fixed compile errors

[aartbik]

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

231

this can go now, right?

240

shall we do the same for index to be consistent?

[K-Wu]

should be minor

[K-Wu]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1261

Seeking ways for easier-to-read literal comparison

[K-Wu]

using the API Wren suggested

[K-Wu]

rebase origin/main

[Peiming]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1261

You should be able to perform the compare without type conversion. It is essentially a pointer comparison and they will share the same underlying storage if they are equal.

[wrengr]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1259–1264

It would be cleaner to just call getConstantIntValue once and save the results. I've attached an edit to show what I mean. (Though really, there's not a whole lot of benefit to the dw_ variable, so you might as well just use *dw directly wherever needed)

Though really this whole pattern should be floated out to a helper function, including the parts that check *dw and return the corresponding FloatType— since the only difference between the version here and the version for mgpuCreateCOO is which op.getOperand is used.

1288

Since gpu::CreateDnMatOp::getMemref returns a TypedValue<MemRefType>, the TypedValue<MemRefType>::getType method will already return a MemRefType, therefore you don't need the llvm::cast

1294–1303

This is yet another place that uses the same function for converting a Value (which has LLVM::ConstantOp defining op) into the appropriate FloatType...

1338–1339

You can replace this with just computeTypeOptional.value_or(defaultType) https://en.cppreference.com/w/cpp/utility/optional/value_or

[aartbik]

can you please add some examples to the ops.mlir for proper roundtrippnig

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1863

ah, good catch on this typo!

1891

I think it is more common in MLIR to put the result type at the end with a :

as in

... : f64

[Peiming]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1256–1257

I don't think you need it (same below for getDnxxx)

[K-Wu]

simpler get element type func

[K-Wu]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1261

I wish I could have known it earlier. Thank you!

[K-Wu]

rm unnecessary MemRefType cast

[K-Wu]

clean up the code a bit

[K-Wu]

rm dw

[K-Wu]

fix to pass tests

[K-Wu]

upd format; add roundtrip test

[K-Wu]

mlir/test/Dialect/GPU/sparse-roundtrip.mlir

42 ↗

(On Diff #526864)

Hi @aartbik , I added a round-trip test where the new syntax : computeType is tested. Let me know what you think. Thanks!

[K-Wu]

refactor index type as well

[Peiming]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1327–1328

No longer needed?

1331

nit: maybe renaming to getConstInt32From?

1335

nit: getConstInt32xxx

[K-Wu]

addressing comments

[K-Wu]

rebase origin/main

[K-Wu]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1327–1328

Good catch! Addressed in new commit

[K-Wu]

Build 364523 failed due to upload failure and is not caused by the code/test itself

[aartbik]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1913

make this Examples

and add an op below that has a compute type set

2000

given the mix of return types and other types, an "into" feels a bit nicer for compute type.
also, do we perhaps simply want to make compute type non-optional and always show it?

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

692

Start with capital C and end with period

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

308

this TODO is done, right?

[K-Wu]

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

308

Good catch. Thanks!

[K-Wu]

addressing comments

[K-Wu]

fix compile errors

[K-Wu]

rebase origin/main

[Peiming]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

695

Here and Below: Do you think it is better to copy the CUSPARSE macro to the file (but probably adding some prefix) so that it can be updated more easily later? Or does the current way already follows the same convention in the file?

[K-Wu]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

695

It is kind of the reverse. Here we map types to integers, but in the cuda header files these enum class definitions are essentially mapping types to an integer, I.e., assigning an integer to each enum item. You raise a good point that we need to have better way to keep in sync, though the CUDA headers seem to keep the mapping of already existing types the same when adding new types probably as a way to keep backward compatibility. I will need to think about how to do that in the future.

[K-Wu]

rebase origin/main

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

let description = [{
 
     The matrix arguments can also be associated with one of the following 
     operators: NON_TRANSPOSE, TRANSPOSE, CONJUGATE_TRANSPOSE. The default value
     is NON_TRANSPOSE.
 
     Example:
 
     ```mlir
-    %buffersz, %token = gpu.spmv_buffersize async [%dep] %env, %spmatA{TRANSPOSE}, %dnX, %dnY
+    %buffersz, %token = gpu.spmv_buffer_size async [%dep] %env, %spmatA{TRANSPOSE}, %dnX, %dnY
     ```
   }];
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                        GPU_SparseEnvHandle:$env,
                        GPU_TransposeModeAttr:$modeA,
                        GPU_SparseSpMatHandle:$spmatA,
                        GPU_SparseDnVecHandle:$dnX,
-                       GPU_SparseDnVecHandle:$dnY);
+                       GPU_SparseDnVecHandle:$dnY,
+                       OptionalAttr<TypeAttr>:$computeType);
   let results = (outs Res<Index>:$bufferSz, 
                       Optional<GPU_AsyncToken>:$asyncToken);
 
   let builders = [OpBuilder<(ins
-      "::mlir::Type":$bufferSz,
-      "::mlir::Type":$asyncToken,
-      "::mlir::ValueRange":$asyncDependencies,
-      "::mlir::Value":$env,
-      "::mlir::Value":$spmatA,
-      "::mlir::Value":$dnX,
-      "::mlir::Value":$dnY), [{
+      "Type":$bufferSz,
+      "Type":$asyncToken,
+      "ValueRange":$asyncDependencies,
+      "Value":$env,
+      "Value":$spmatA,
+      "Value":$dnX,
+      "Value":$dnY)
+      , [{
     auto modeA = gpu::TransposeMode::NON_TRANSPOSE;
-    return build($_builder, $_state, bufferSz, asyncToken, asyncDependencies, env, 
-                 modeA, spmatA, dnX, dnY);}]>
+    return build($_builder, $_state, bufferSz, asyncToken, asyncDependencies, 
+                 env, modeA, spmatA, dnX, dnY, {});}]>
   ];
 
   let assemblyFormat = [{
+    (`{` $computeType^ `}`)?
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnX `,` $dnY attr-dict
   }];
 }
 
 def GPU_SpMVOp : GPU_Op<"spmv", [GPU_AsyncOpInterface]> {
   let summary = "SpMV operation";
   let description = [{
     The `gpu.spmv` operation performs the SpMV operation on the given sparse matrix,
     dense vectors, and buffer.  The operation expects handles returned by previous
     sparse operations to construct an environment and the operands for SpMV. The
     buffer must have been allocated on the device.
 
     If the `async` keyword is present, the op is executed asynchronously (i.e.
     it does not block until the execution has finished on the device). In
     that case, it returns a !gpu.async.token in addition to the environment.
 
     The matrix arguments can also be associated with one of the following 
     operators: NON_TRANSPOSE, TRANSPOSE, CONJUGATE_TRANSPOSE. The default value
     is NON_TRANSPOSE.
 
     Example:
 
     ```mlir
     %token = gpu.spmv async [%dep] %env, %spmatA{TRANSPOSE}, %dnX, %dnY : memref<?xf64>
     ```
   }];
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                        GPU_SparseEnvHandle:$env,
                        GPU_TransposeModeAttr:$modeA,
                        GPU_SparseSpMatHandle:$spmatA,
                        GPU_SparseDnVecHandle:$dnX,
                        GPU_SparseDnVecHandle:$dnY,
+                       OptionalAttr<TypeAttr>:$computeType,
                        AnyMemRef:$buffer);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let builders = [OpBuilder<(ins
-      "::mlir::Type":$asyncToken,
-      "::mlir::ValueRange":$asyncDependencies,
-      "::mlir::Value":$env,
-      "::mlir::Value":$spmatA,
-      "::mlir::Value":$dnX,
-      "::mlir::Value":$dnY,
-      "::mlir::Value":$buffer), [{
+      "Type":$asyncToken,
+      "ValueRange":$asyncDependencies,
+      "Value":$env,
+      "Value":$spmatA,
+      "Value":$dnX,
+      "Value":$dnY,
+      "Value":$buffer), [{
     auto modeA = gpu::TransposeMode::NON_TRANSPOSE;
     return build($_builder, $_state, asyncToken, asyncDependencies, env, modeA,
-                 spmatA, dnX, dnY, buffer);}]>
+                 spmatA, dnX, dnY, {}, buffer);}]>
   ];
 
   let assemblyFormat = [{
+    (`{` $computeType^ `}`)?
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnX `,` $dnY `,` $buffer attr-dict `:` type($buffer)
   }];
 }
 
 def GPU_SpMMBufferSizeOp : GPU_Op<"spmm_buffer_size", [GPU_AsyncOpInterface]> {
   let summary = "Precompute buffersize for SpMM operation";
   let description = [{
def GPU_SpMMBufferSizeOp : GPU_Op<"spmm_buffer_size", [GPU_AsyncOpInterface]> {
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                        GPU_SparseEnvHandle:$env,
                        GPU_TransposeModeAttr:$modeA,
                        GPU_TransposeModeAttr:$modeB,
                        GPU_SparseSpMatHandle:$spmatA,
                        GPU_SparseDnMatHandle:$dnmatB,
-                       GPU_SparseDnMatHandle:$dnmatC);
+                       GPU_SparseDnMatHandle:$dnmatC,
+                       OptionalAttr<TypeAttr>:$computeType);
   let results = (outs Res<Index>:$bufferSz, 
                       Optional<GPU_AsyncToken>:$asyncToken);
 
   let builders = [OpBuilder<(ins
-      "::mlir::Type":$bufferSz,
-      "::mlir::Type":$asyncToken,
-      "::mlir::ValueRange":$asyncDependencies,
-      "::mlir::Value":$env,
-      "::mlir::Value":$spmatA,
-      "::mlir::Value":$dnmatB,
-      "::mlir::Value":$dnmatC), [{
+      "Type":$bufferSz,
+      "Type":$asyncToken,
+      "ValueRange":$asyncDependencies,
+      "Value":$env,
+      "Value":$spmatA,
+      "Value":$dnmatB,
+      "Value":$dnmatC), [{
     auto modeA = gpu::TransposeMode::NON_TRANSPOSE;
     auto modeB = gpu::TransposeMode::NON_TRANSPOSE;
     return build($_builder, $_state, bufferSz, asyncToken, asyncDependencies, 
-                 env, modeA, modeB, spmatA, dnmatB, dnmatC);}]>
+                 env, modeA, modeB, spmatA, dnmatB, dnmatC, {});}]>
   ];
 
   let assemblyFormat = [{
+    (`{` $computeType^ `}`)?
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $dnmatC attr-dict
   }];
 }
 
 def GPU_SpMMOp : GPU_Op<"spmm", [GPU_AsyncOpInterface]> {
   let summary = "SpMM operation";
   let description = [{
     The `gpu.spmm` operation performs the SpMM operation on the given sparse and
     dense matrix, and buffer.  The operation expects handles returned by previous
def GPU_SpMMOp : GPU_Op<"spmm", [GPU_AsyncOpInterface]> {
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                        GPU_SparseEnvHandle:$env,
                        GPU_TransposeModeAttr:$modeA,
                        GPU_TransposeModeAttr:$modeB,
                        GPU_SparseSpMatHandle:$spmatA,
                        GPU_SparseDnMatHandle:$dnmatB,
                        GPU_SparseDnMatHandle:$dnmatC,
+                       OptionalAttr<TypeAttr>:$computeType,
                        AnyMemRef:$buffer);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let builders = [OpBuilder<(ins
-      "::mlir::Type":$asyncToken,
-      "::mlir::ValueRange":$asyncDependencies,
-      "::mlir::Value":$env,
-      "::mlir::Value":$spmatA,
-      "::mlir::Value":$dnmatB,
-      "::mlir::Value":$dnmatC,
-      "::mlir::Value":$buffer), [{
+      "Type":$asyncToken,
+      "ValueRange":$asyncDependencies,
+      "Value":$env,
+      "Value":$spmatA,
+      "Value":$dnmatB,
+      "Value":$dnmatC,
+      "Value":$buffer), [{
     auto modeA = gpu::TransposeMode::NON_TRANSPOSE;
     auto modeB = gpu::TransposeMode::NON_TRANSPOSE;
     return build($_builder, $_state, asyncToken, asyncDependencies, env, modeA, 
-                 modeB, spmatA, dnmatB, dnmatC, buffer);}]>
+                 modeB, spmatA, dnmatB, dnmatC, {}, buffer);}]>
   ];
 
   let assemblyFormat = [{
+    (`{` $computeType^ `}`)?
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $dnmatC `,` $buffer attr-dict `:` type($buffer)
   }];
 }
 
 def GPU_SDDMMBufferSizeOp : GPU_Op<"sddmm_buffer_size", [GPU_AsyncOpInterface]> {
   let summary = "Precompute buffersize for SDDMM operation";
   let description = [{
def GPU_SDDMMBufferSizeOp : GPU_Op<"sddmm_buffer_size", [GPU_AsyncOpInterface]> {
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                    GPU_SparseEnvHandle:$env,
                    GPU_TransposeModeAttr:$modeA,
                    GPU_TransposeModeAttr:$modeB,
                    GPU_SparseDnMatHandle:$dnmatA,
                    GPU_SparseDnMatHandle:$dnmatB,
-                   GPU_SparseSpMatHandle:$spmatC);
+                   GPU_SparseSpMatHandle:$spmatC,
+                   OptionalAttr<TypeAttr>:$computeType);
   let results = (outs Res<Index>:$bufferSz, Optional<GPU_AsyncToken>:$asyncToken);
 
   let builders = [OpBuilder<(ins
       "::mlir::Type":$bufferSz,
       "::mlir::Type":$asyncToken,
       "::mlir::ValueRange":$asyncDependencies,
       "::mlir::Value":$env,
       "::mlir::Value":$dnmatA,
       "::mlir::Value":$dnmatB,
       "::mlir::Value":$spmatC), [{
     auto modeA = gpu::TransposeMode::NON_TRANSPOSE;
     auto modeB = gpu::TransposeMode::NON_TRANSPOSE;
     return build($_builder, $_state, bufferSz, asyncToken, asyncDependencies, 
-                 env, modeA, modeB, dnmatA, dnmatB, spmatC);}]>
+                 env, modeA, modeB, dnmatA, dnmatB, spmatC, {});}]>
   ];
 
   let assemblyFormat = [{
+    (`{` $computeType^ `}`)?
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $env `,` $dnmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $spmatC attr-dict
   }];
 }
 
 def GPU_SDDMMOp : GPU_Op<"sddmm", [GPU_AsyncOpInterface]> {
   let summary = "SDDMM operation";
   let description = [{
def GPU_SDDMMOp : GPU_Op<"sddmm", [GPU_AsyncOpInterface]> {
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                    GPU_SparseEnvHandle:$env,
                    GPU_TransposeModeAttr:$modeA,
                    GPU_TransposeModeAttr:$modeB,
                    GPU_SparseDnMatHandle:$dnmatA,
                    GPU_SparseDnMatHandle:$dnmatB,
                    GPU_SparseSpMatHandle:$spmatC,
+                   OptionalAttr<TypeAttr>:$computeType,
                    AnyMemRef:$buffer);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let builders = [OpBuilder<(ins
     "::mlir::Type":$asyncToken,
     "::mlir::ValueRange":$asyncDependencies,
     "::mlir::Value":$env,
     "::mlir::Value":$dnmatA,
     "::mlir::Value":$dnmatB,
     "::mlir::Value":$spmatC,
     "::mlir::Value":$buffer), [{
   auto modeA = gpu::TransposeMode::NON_TRANSPOSE;
   auto modeB = gpu::TransposeMode::NON_TRANSPOSE;
   return build($_builder, $_state, asyncToken, asyncDependencies, env, modeA, 
-                modeB, dnmatA, dnmatB, spmatC, buffer);}]>
+                modeB, dnmatA, dnmatB, spmatC, {}, buffer);}]>
   ];
 
   let assemblyFormat = [{
+    (`{` $computeType^ `}`)?
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $env `,` $dnmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $spmatC `,` $buffer attr-dict `:` type($buffer)
   }];
 }
 
 #endif // GPU_OPS

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

protected:
   FunctionCallBuilder destroySpMatCallBuilder = {
       "mgpuDestroySpMat",
       llvmVoidType,
       {llvmPointerType, llvmPointerType /* void *stream */}};
   FunctionCallBuilder spMVBufferSizeCallBuilder = {
       "mgpuSpMVBufferSize",
       llvmIntPtrType,
       {llvmPointerType, llvmInt32Type, llvmPointerType, llvmPointerType,
-       llvmPointerType, llvmInt32Type, llvmPointerType /* void *stream */}};
+       llvmPointerType, llvmInt32Type, llvmInt32Type,
+       llvmPointerType /* void *stream */}};
   FunctionCallBuilder spMVCallBuilder = {
       "mgpuSpMV",
       llvmVoidType,
       {llvmPointerType, llvmInt32Type, llvmPointerType, llvmPointerType,
-       llvmPointerType, llvmInt32Type, llvmPointerType,
+       llvmPointerType, llvmInt32Type, llvmInt32Type, llvmPointerType,
        llvmPointerType /* void *stream */}};
   FunctionCallBuilder spMMBufferSizeCallBuilder = {
       "mgpuSpMMBufferSize",
       llvmIntPtrType,
       {llvmPointerType, llvmInt32Type, llvmInt32Type, llvmPointerType,
-       llvmPointerType, llvmPointerType, llvmInt32Type,
+       llvmPointerType, llvmPointerType, llvmInt32Type, llvmInt32Type,
        llvmPointerType /* void *stream */}};
   FunctionCallBuilder spMMCallBuilder = {
       "mgpuSpMM",
       llvmVoidType,
       {llvmPointerType, llvmInt32Type, llvmInt32Type, llvmPointerType,
-       llvmPointerType, llvmPointerType, llvmInt32Type, llvmPointerType,
-       llvmPointerType /* void *stream */}};
+       llvmPointerType, llvmPointerType, llvmInt32Type, llvmInt32Type,
+       llvmPointerType, llvmPointerType /* void *stream */}};
   FunctionCallBuilder SDDMMBufferSizeCallBuilder = {
       "mgpuSDDMMBufferSize",
       llvmIntPtrType,
       {llvmPointerType, llvmInt32Type, llvmInt32Type, llvmPointerType,
-       llvmPointerType, llvmPointerType, llvmInt32Type,
+       llvmPointerType, llvmPointerType, llvmInt32Type, llvmInt32Type,
        llvmPointerType /* void *stream */}};
   FunctionCallBuilder SDDMMCallBuilder = {
       "mgpuSDDMM",
       llvmVoidType,
       {llvmPointerType, llvmInt32Type, llvmInt32Type, llvmPointerType,
-       llvmPointerType, llvmPointerType, llvmInt32Type, llvmPointerType,
-       llvmPointerType /* void *stream */}};
+       llvmPointerType, llvmPointerType, llvmInt32Type, llvmInt32Type,
+       llvmPointerType, llvmPointerType /* void *stream */}};
 };
 
 /// A rewrite pattern to convert gpu.host_register operations into a GPU runtime
 /// call. Currently it supports CUDA and ROCm (HIP).
 class ConvertHostRegisterOpToGpuRuntimeCallPattern
     : public ConvertOpToGpuRuntimeCallPattern<gpu::HostRegisterOp> {
 public:
   ConvertHostRegisterOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter)
auto function = [&] {
     if (auto function = module.lookupSymbol<LLVM::LLVMFuncOp>(functionName))
       return function;
     return OpBuilder::atBlockEnd(module.getBody())
         .create<LLVM::LLVMFuncOp>(loc, functionName, functionType);
   }();
   return builder.create<LLVM::CallOp>(loc, function, arguments);
 }
 
+// corresponding to cudaDataType_t defined in library_types.h
+// TODO: add support to complex types
+static int32_t getCUSparseDataTypeEnumFrom(Type type) {
+  if (type.isBF16())
+    return 14;
+  if (type.isF16())
+    return 2;
+  if (type.isF32())
+    return 0;
+  if (type.isF64())
+    return 1;
+  if (type.isInteger(8))
+    return 3;
+  if (type.isInteger(32))
+    return 10;
+  llvm_unreachable("unsupported element type");
+}
+
 // Returns whether all operands are of LLVM type.
 static LogicalResult areAllLLVMTypes(Operation *op, ValueRange operands,
                                      ConversionPatternRewriter &rewriter) {
   if (!llvm::all_of(operands, [](Value value) {
         return LLVM::isCompatibleType(value.getType());
       }))
     return rewriter.notifyMatchFailure(
         op, "Cannot convert if operands aren't of LLVM type.");

 
 // Returns the element type of the defining spmat op.
 // TODO: safer and more flexible to store data type in actual op instead?
 static Type getSpMatElemType(Value spMat) {
   if (auto op = spMat.getDefiningOp<gpu::CreateCooOp>())
     return llvm::cast<MemRefType>(op.getValues().getType()).getElementType();
   if (auto op = spMat.getDefiningOp<gpu::CreateCsrOp>())
     return llvm::cast<MemRefType>(op.getValues().getType()).getElementType();
+  // the defining op may also be llvm.call after partial lowering
+  if (auto op = spMat.getDefiningOp<LLVM::CallOp>()) {
+    if (op.getCallee() == "mgpuCreateCsr") {
+      mlir::Attribute dw =
+          op.getOperand(8).getDefiningOp<LLVM::ConstantOp>().getValue();
+      if (!getConstantIntValue(dw).has_value()) {
+        llvm_unreachable("expecting dw to be a constant but not");
+      }
+      auto dw_ = getConstantIntValue(dw).value();
+      if (dw_ == 32)
+        return FloatType::getF32(spMat.getContext());
+      else if (dw_ == 64)
+        return FloatType::getF64(spMat.getContext());
+    } else if (op.getCallee() == "mgpuCreateCoo") {
+      mlir::Attribute dw =
+          op.getOperand(7).getDefiningOp<LLVM::ConstantOp>().getValue();
+      if (!getConstantIntValue(dw).has_value()) {
+        llvm_unreachable("expecting dw to be a constant but not");
+      }
+      auto dw_ = getConstantIntValue(dw).value();
+      if (dw_ == 32)
+        return FloatType::getF32(spMat.getContext());
+      else if (dw_ == 64)
+        return FloatType::getF64(spMat.getContext());
+    }
+  }
   llvm_unreachable("cannot find spmat def");
 }
 
+// Returns the element type of the defining dnmat or dnvec op.
+static Type getDnElemType(Value dn) {
+  if (auto op = dn.getDefiningOp<gpu::CreateDnMatOp>())
+    return llvm::cast<MemRefType>(op.getMemref().getType()).getElementType();
+  if (auto op = dn.getDefiningOp<gpu::CreateDnVecOp>())
+    return llvm::cast<MemRefType>(op.getMemref().getType()).getElementType();
+  // the defining op may also be llvm.call after partial lowering
+  if (auto op = dn.getDefiningOp<LLVM::CallOp>()) {
+    if (op.getCallee() == "mgpuCreateDnVec") {
+      mlir::Attribute dw =
+          op.getOperand(2).getDefiningOp<LLVM::ConstantOp>().getValue();
+      if (!getConstantIntValue(dw).has_value()) {
+        llvm_unreachable("expecting dw to be a constant but not");
+      }
+      auto dw_ = getConstantIntValue(dw).value();
+      if (dw_ == 32)
+        return FloatType::getF32(dn.getContext());
+      else if (dw_ == 64)
+        return FloatType::getF64(dn.getContext());
+    } else if (op.getCallee() == "mgpuCreateDnMat") {
+      mlir::Attribute dw =
+          op.getOperand(3).getDefiningOp<LLVM::ConstantOp>().getValue();
+      if (!getConstantIntValue(dw).has_value()) {
+        llvm_unreachable("expecting dw to be a constant but not");
+      }
+      auto dw_ = getConstantIntValue(dw).value();
+      if (dw_ == 32)
+        return FloatType::getF32(dn.getContext());
+      else if (dw_ == 64)
+        return FloatType::getF64(dn.getContext());
+    }
+  }
+  llvm_unreachable("cannot find dn def");
+}
+
 static Value genConstFrom(OpBuilder &builder, Location loc,
                           gpu::TransposeMode mode) {
   Type llvmInt32Type = builder.getIntegerType(32);
   return builder.create<LLVM::ConstantOp>(loc, llvmInt32Type,
                                           static_cast<int32_t>(mode));
 }
 
+static Value genConstFrom(OpBuilder &builder, Location loc,
+                          int32_t computeTypeInt) {
+  Type llvmInt32Type = builder.getIntegerType(32);
+  return builder.create<LLVM::ConstantOp>(loc, llvmInt32Type, computeTypeInt);
+}
+
+static Value
+genConstFromOptionalComputeMode(OpBuilder &builder, Location loc,
+                                std::optional<Type> computeTypeOptional,
+                                Type defaultType) {
+  auto computeTypeInt = getCUSparseDataTypeEnumFrom(
+      computeTypeOptional.has_value() ? computeTypeOptional.value()
+                                      : defaultType);
+  auto computeType = genConstFrom(builder, loc, computeTypeInt);
+  return computeType;
+}
+
 LogicalResult ConvertCreateSparseEnvOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::CreateSparseEnvOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
   auto stream = adaptor.getAsyncDependencies().front();
LogicalResult ConvertSpMVBufferSizeOpToGpuRuntimeCallPattern::matchAndRewrite(
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
   auto modeA = genConstFrom(rewriter, loc, op.getModeA());
   Type dType = getSpMatElemType(op.getSpmatA());
   auto dw = rewriter.create<LLVM::ConstantOp>(loc, llvmInt32Type,
                                               dType.getIntOrFloatBitWidth());
+  // retrieve the compute type, notice that it may be optional
+  auto computeType = genConstFromOptionalComputeMode(
+      rewriter, loc, adaptor.getComputeType(), getDnElemType(adaptor.getDnY()));
   auto stream = adaptor.getAsyncDependencies().front();
   auto bufferSize =
       spMVBufferSizeCallBuilder
           .create(loc, rewriter,
                   {adaptor.getEnv(), modeA, adaptor.getSpmatA(),
-                   adaptor.getDnX(), adaptor.getDnY(), dw, stream})
+                   adaptor.getDnX(), adaptor.getDnY(), dw, computeType, stream})
           .getResult();
   rewriter.replaceOp(op, {bufferSize, stream});
   return success();
 }
 
 LogicalResult ConvertSpMVOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::SpMVOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
   Type dType = getSpMatElemType(op.getSpmatA());
   auto modeA = genConstFrom(rewriter, loc, adaptor.getModeA());
   auto dw = rewriter.create<LLVM::ConstantOp>(loc, llvmInt32Type,
                                               dType.getIntOrFloatBitWidth());
+  // retrieve the compute type, notice that it may be optional
+  auto computeType = genConstFromOptionalComputeMode(
+      rewriter, loc, adaptor.getComputeType(), getDnElemType(adaptor.getDnY()));
   auto stream = adaptor.getAsyncDependencies().front();
   Value pBuf =
       MemRefDescriptor(adaptor.getBuffer()).allocatedPtr(rewriter, loc);
   if (!getTypeConverter()->useOpaquePointers())
     pBuf = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, pBuf);
   spMVCallBuilder.create(loc, rewriter,
                          {adaptor.getEnv(), modeA, adaptor.getSpmatA(),
-                          adaptor.getDnX(), adaptor.getDnY(), dw, pBuf,
-                          stream});
+                          adaptor.getDnX(), adaptor.getDnY(), dw, computeType,
+                          pBuf, stream});
   rewriter.replaceOp(op, {stream});
   return success();
 }
 
 LogicalResult ConvertSpMMBufferSizeOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::SpMMBufferSizeOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
   auto modeA = genConstFrom(rewriter, loc, adaptor.getModeA());
   auto modeB = genConstFrom(rewriter, loc, adaptor.getModeB());
   Type dType = getSpMatElemType(op.getSpmatA());
   auto dw = rewriter.create<LLVM::ConstantOp>(loc, llvmInt32Type,
                                               dType.getIntOrFloatBitWidth());
   auto stream = adaptor.getAsyncDependencies().front();
-  auto bufferSize =
-      spMMBufferSizeCallBuilder
+  // retrieve the compute type, notice that it may be optional
+  auto computeType =
+      genConstFromOptionalComputeMode(rewriter, loc, adaptor.getComputeType(),
+                                      getDnElemType(adaptor.getDnmatC()));
+
+  auto bufferSize = spMMBufferSizeCallBuilder
                         .create(loc, rewriter,
-                  {adaptor.getEnv(), modeA, modeB, adaptor.getSpmatA(),
-                   adaptor.getDnmatB(), adaptor.getDnmatC(), dw, stream})
+                                {adaptor.getEnv(), modeA, modeB,
+                                 adaptor.getSpmatA(), adaptor.getDnmatB(),
+                                 adaptor.getDnmatC(), dw, computeType, stream})
                         .getResult();
   rewriter.replaceOp(op, {bufferSize, stream});
   return success();
 }
 
 LogicalResult ConvertSDDMMBufferSizeOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::SDDMMBufferSizeOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
   auto modeA = genConstFrom(rewriter, loc, adaptor.getModeA());
   auto modeB = genConstFrom(rewriter, loc, adaptor.getModeB());
   Type dType = getSpMatElemType(op.getSpmatC());
   auto dw = rewriter.create<LLVM::ConstantOp>(loc, llvmInt32Type,
                                               dType.getIntOrFloatBitWidth());
+  auto computeType =
+      genConstFromOptionalComputeMode(rewriter, loc, adaptor.getComputeType(),
+                                      getSpMatElemType(adaptor.getSpmatC()));
   auto stream = adaptor.getAsyncDependencies().front();
-  auto bufferSize =
-      SDDMMBufferSizeCallBuilder
+  auto bufferSize = SDDMMBufferSizeCallBuilder
                         .create(loc, rewriter,
-                  {adaptor.getEnv(), modeA, modeB, adaptor.getDnmatA(),
-                   adaptor.getDnmatB(), adaptor.getSpmatC(), dw, stream})
+                                {adaptor.getEnv(), modeA, modeB,
+                                 adaptor.getDnmatA(), adaptor.getDnmatB(),
+                                 adaptor.getSpmatC(), dw, computeType, stream})
                         .getResult();
   rewriter.replaceOp(op, {bufferSize, stream});
   return success();
 }
 
 LogicalResult ConvertSpMMOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::SpMMOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
   auto modeA = genConstFrom(rewriter, loc, adaptor.getModeA());
   auto modeB = genConstFrom(rewriter, loc, adaptor.getModeB());
   Type dType = getSpMatElemType(op.getSpmatA());
   auto dw = rewriter.create<LLVM::ConstantOp>(loc, llvmInt32Type,
                                               dType.getIntOrFloatBitWidth());
+  // retrieve the compute type, notice that it may be optional
+  auto computeType =
+      genConstFromOptionalComputeMode(rewriter, loc, adaptor.getComputeType(),
+                                      getDnElemType(adaptor.getDnmatC()));
+
   auto stream = adaptor.getAsyncDependencies().front();
   Value pBuf =
       MemRefDescriptor(adaptor.getBuffer()).allocatedPtr(rewriter, loc);
   if (!getTypeConverter()->useOpaquePointers())
     pBuf = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, pBuf);
   spMMCallBuilder.create(loc, rewriter,
                          {adaptor.getEnv(), modeA, modeB, adaptor.getSpmatA(),
-                          adaptor.getDnmatB(), adaptor.getDnmatC(), dw, pBuf,
-                          stream});
+                          adaptor.getDnmatB(), adaptor.getDnmatC(), dw,
+                          computeType, pBuf, stream});
   rewriter.replaceOp(op, {stream});
   return success();
 }
 
 template <typename T>
 static void addOpaquePointerConversion(LLVMTypeConverter &converter) {
   converter.addConversion([&converter](T) -> Type {
     return converter.getPointerType(
         IntegerType::get(&converter.getContext(), 8));
   });
 }
 
 LogicalResult ConvertSDDMMOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::SDDMMOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
   Type dType = getSpMatElemType(op.getSpmatC());
   auto dw = rewriter.create<LLVM::ConstantOp>(loc, llvmInt32Type,
                                               dType.getIntOrFloatBitWidth());
+  auto computeType =
+      genConstFromOptionalComputeMode(rewriter, loc, adaptor.getComputeType(),
+                                      getSpMatElemType(adaptor.getSpmatC()));
   auto modeA = genConstFrom(rewriter, loc, adaptor.getModeA());
   auto modeB = genConstFrom(rewriter, loc, adaptor.getModeB());
   auto stream = adaptor.getAsyncDependencies().front();
   Value pBuf =
       MemRefDescriptor(adaptor.getBuffer()).allocatedPtr(rewriter, loc);
   if (!getTypeConverter()->useOpaquePointers())
     pBuf = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, pBuf);
   SDDMMCallBuilder.create(loc, rewriter,
                           {adaptor.getEnv(), modeA, modeB, adaptor.getDnmatA(),
-                           adaptor.getDnmatB(), adaptor.getSpmatC(), dw, pBuf,
-                           stream});
+                           adaptor.getDnmatB(), adaptor.getSpmatC(), dw,
+                           computeType, pBuf, stream});
   rewriter.replaceOp(op, {stream});
   return success();
 }
 
 void mlir::populateGpuToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                                RewritePatternSet &patterns,
                                                StringRef gpuBinaryAnnotation,
                                                bool kernelBarePtrCallConv) {

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSetDefaultDevice(int32_t device) {
   defaultDevice = device;
 }
 
 ///
 /// Wrapper methods for the cuSparse library.
 ///
 
 static inline cudaDataType_t dataTp(int32_t width) {
   switch (width) {
   case 32:
     return CUDA_R_32F;
   default:
     return CUDA_R_64F;
   }
 }
 
 static inline cusparseIndexType_t idxTp(int32_t width) {
   switch (width) {
   case 32:
     return CUSPARSE_INDEX_32I;
   default:
     return CUSPARSE_INDEX_64I;
   }
 }
 

 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
 mgpuDestroyDnMat(void *m, CUstream /*stream*/) {
   cusparseDnMatDescr_t mat = reinterpret_cast<cusparseDnMatDescr_t>(m);
   CUSPARSE_REPORT_IF_ERROR(cusparseDestroyDnMat(mat))
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void *
 mgpuCreateCoo(intptr_t rows, intptr_t cols, intptr_t nnz, void *rowIdxs,
               void *colIdxs, void *values, int32_t iw, int32_t dw,
               CUstream /*stream*/) {
   cusparseSpMatDescr_t mat = nullptr;
   cusparseIndexType_t itp = idxTp(iw);
   cudaDataType_t dtp = dataTp(dw);
   CUSPARSE_REPORT_IF_ERROR(cusparseCreateCoo(&mat, rows, cols, nnz, rowIdxs,
                                              colIdxs, values, itp,

 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
 mgpuDestroySpMat(void *m, CUstream /*stream*/) {
   cusparseSpMatDescr_t mat = reinterpret_cast<cusparseSpMatDescr_t>(m);
   CUSPARSE_REPORT_IF_ERROR(cusparseDestroySpMat(mat))
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t
 mgpuSpMVBufferSize(void *h, int32_t ma, void *a, void *x, void *y, int32_t dw,
-                   CUstream /*stream*/) {
+                   int32_t dtp, CUstream /*stream*/) {
   cusparseHandle_t handle = reinterpret_cast<cusparseHandle_t>(h);
   cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma);
   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);
+  cudaDataType_t dTp = static_cast<cudaDataType_t>(dtp);
   ALPHABETA(dw, alpha, beta)
   size_t bufferSize = 0;
   CUSPARSE_REPORT_IF_ERROR(
       cusparseSpMV_bufferSize(handle, modeA, alphap, matA, vecX, betap, vecY,
-                              dtp, CUSPARSE_SPMV_ALG_DEFAULT, &bufferSize))
+                              dTp, CUSPARSE_SPMV_ALG_DEFAULT, &bufferSize))
   return bufferSize == 0 ? 1 : bufferSize; // avoid zero-alloc
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSpMV(void *h, int32_t ma, void *a,
                                                    void *x, void *y, int32_t dw,
-                                                   void *buf,
+                                                   int32_t dtp, void *buf,
                                                    CUstream /*stream*/) {
   cusparseHandle_t handle = reinterpret_cast<cusparseHandle_t>(h);
   cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma);
   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);
+  cudaDataType_t dTp = static_cast<cudaDataType_t>(dtp);
   ALPHABETA(dw, alpha, beta)
   CUSPARSE_REPORT_IF_ERROR(cusparseSpMV(handle, modeA, alphap, matA, vecX,
-                                        betap, vecY, dtp,
+                                        betap, vecY, dTp,
                                         CUSPARSE_SPMV_ALG_DEFAULT, buf))
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t
 mgpuSpMMBufferSize(void *h, int32_t ma, int32_t mb, void *a, void *b, void *c,
-                   int32_t dw, CUstream /*stream*/) {
+                   int32_t dw, int32_t dtp, CUstream /*stream*/) {
   cusparseHandle_t handle = reinterpret_cast<cusparseHandle_t>(h);
   cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma);
   cusparseOperation_t modeB = static_cast<cusparseOperation_t>(mb);
   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);
+  cudaDataType_t dTp = static_cast<cudaDataType_t>(dtp);
   ALPHABETA(dw, alpha, beta)
   size_t bufferSize = 0;
   CUSPARSE_REPORT_IF_ERROR(cusparseSpMM_bufferSize(
-      handle, modeA, modeB, alphap, matA, matB, betap, matC, dtp,
+      handle, modeA, modeB, 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, int32_t ma, int32_t mb, void *a, void *b, void *c, int32_t dw,
-         void *buf, CUstream /*stream*/) {
+         int32_t dtp, void *buf, CUstream /*stream*/) {
   cusparseHandle_t handle = reinterpret_cast<cusparseHandle_t>(h);
   cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma);
   cusparseOperation_t modeB = static_cast<cusparseOperation_t>(mb);
   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);
+  cudaDataType_t dTp = static_cast<cudaDataType_t>(dtp);
   ALPHABETA(dw, alpha, beta)
   CUSPARSE_REPORT_IF_ERROR(cusparseSpMM(handle, modeA, modeB, alphap, matA,
-                                        matB, betap, matC, dtp,
+                                        matB, betap, matC, dTp,
                                         CUSPARSE_SPMM_ALG_DEFAULT, buf))
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t
 mgpuSDDMMBufferSize(void *h, int32_t ma, int32_t mb, void *a, void *b, void *c,
                     int32_t dw, CUstream /*stream*/) {
   cusparseHandle_t handle = reinterpret_cast<cusparseHandle_t>(h);
   cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma);