[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

232–233

this can go now, right?

232–233

shall we do the same for index to be consistent?

[K-Wu]

should be minor

[K-Wu]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1289

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

1289

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

1287–1292

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.

1290

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

1296–1305

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

1331–1332

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

1284–1285

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

1289

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

43

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

1297

nit: maybe renaming to getConstInt32From?

1303–1304

No longer needed?

1311

nit: getConstInt32xxx

[K-Wu]

addressing comments

[K-Wu]

rebase origin/main

[K-Wu]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1303–1304

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

1912

make this Examples

and add an op below that has a compute type set

1997

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

691

Start with capital C and end with period

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

304

this TODO is done, right?

[K-Wu]

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

304

Good catch. Thanks!

[K-Wu]

addressing comments

[K-Wu]

fix compile errors

[K-Wu]

rebase origin/main

[Peiming]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

694

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

694

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 = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnX `,` $dnY attr-dict
+    $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnX `,` $dnY attr-dict ( `:` $computeType^)?
   }];
 }
 
 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 = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnX `,` $dnY `,` $buffer attr-dict `:` type($buffer)
+    $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnX `,` $dnY `,` $buffer attr-dict `:` type($buffer) ( `,` $computeType^)?
   }];
 }
 
 def GPU_SpMMBufferSizeOp : GPU_Op<"spmm_buffer_size", [GPU_AsyncOpInterface]> {
   let summary = "Precompute buffersize for SpMM operation";
   let description = [{
     The `gpu.spmm_buffer_size` operation returns the buffer size required
     to perform the SpMM operation on the given sparse and dense matrix.
     The operation expects handles returned by previous sparse operations
     to construct an environment and the operands for SpMM.
 
     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
-    %buffersz, %token = gpu.spmm_buffersize async [%dep] %env, %spmatA{TRANSPOSE}, %dnmatB{TRANSPOSE}, %dnmatC
+    %buffersz, %token = gpu.spmm_buffer_size async [%dep] %env, %spmatA{TRANSPOSE}, %dnmatB{TRANSPOSE}, %dnmatC
     ```
   }];
 
   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 = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $dnmatC attr-dict
+    $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $dnmatC attr-dict ( `:` $computeType^)?
   }];
 }
 
 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 = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $dnmatC `,` $buffer attr-dict `:` type($buffer)
+    $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $dnmatC `,` $buffer attr-dict `:` type($buffer) ( `,` $computeType^)?
   }];
 }
 
 def GPU_SDDMMBufferSizeOp : GPU_Op<"sddmm_buffer_size", [GPU_AsyncOpInterface]> {
   let summary = "Precompute buffersize for SDDMM operation";
   let description = [{
     The `gpu.sddmm_buffer_size` operation returns the buffer size required
     to perform the SDDMM operation on the given sparse and dense matrices.
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 = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $env `,` $dnmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $spmatC attr-dict
+    $env `,` $dnmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $spmatC attr-dict ( `:` $computeType^)?
   }];
 }
 
 def GPU_SDDMMOp : GPU_Op<"sddmm", [GPU_AsyncOpInterface]> {
   let summary = "SDDMM operation";
   let description = [{
     The `gpu.sddmm` operation performs the SDDMM operation on the given sparse and
     dense matrices, and buffer.  The operation expects handles returned by previous
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 = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $env `,` $dnmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $spmatC `,` $buffer attr-dict `:` type($buffer)
+    $env `,` $dnmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $spmatC `,` $buffer attr-dict `:` type($buffer) ( `,` $computeType^)?
   }];
 }
 
 #endif // GPU_OPS

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

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 cusparseIndexType_t defined in cusparse.h
+static int32_t getCuSparseIndexTypeFrom(Type type) {
+  if (type.isa<IndexType>())
+    return 3; // CUSPARSE_INDEX_64I
+  else
+    return 2; // CUSPARSE_INDEX_32I
+  // TODO: add support to CUSPARSE_INDEX_16U: 1
+}
+
+// corresponding to cudaDataType_t defined in CUDA library_types.h
+static int32_t getCuSparseDataTypeFrom(Type type) {
+  if (llvm::isa<ComplexType>(type)) {
+    // get the element type
+    auto elementType = type.cast<ComplexType>().getElementType();
+    if (elementType.isBF16())
+      return 15; // CUDA_C_16BF
+    if (elementType.isF16())
+      return 6; // CUDA_C_16F
+    if (elementType.isF32())
+      return 4; // CUDA_C_32F
+    if (elementType.isF64())
+      return 5; // CUDA_C_64F
+    if (elementType.isInteger(8))
+      return 7; // CUDA_C_8I
+    if (elementType.isInteger(16))
+      return 21; // CUDA_C_16I
+    if (elementType.isInteger(32))
+      return 11; // CUDA_C_32I
+  }
+  if (type.isBF16())
+    return 14; // CUDA_R_16BF
+  if (type.isF16())
+    return 2; // CUDA_R_16F
+  if (type.isF32())
+    return 0; // CUDA_R_32F
+  if (type.isF64())
+    return 1; // CUDA_R_64F
+  if (type.isInteger(8))
+    return 3; // CUDA_R_8I
+  if (type.isInteger(16))
+    return 20; // CUDA_R_16I
+  if (type.isInteger(32))
+    return 10; // CUDA_R_32I
+
+  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();
   llvm_unreachable("cannot find spmat def");
 }
 
-static Value genConstFrom(OpBuilder &builder, Location loc,
-                          gpu::TransposeMode mode) {
+// Returns the element type of the defining dnmat or dnvec op.
+static Type getDnElemType(Value dn) {
+  if (auto op = dn.getDefiningOp<gpu::CreateDnMatOp>())
+    return op.getMemref().getType().getElementType();
+  if (auto op = dn.getDefiningOp<gpu::CreateDnVecOp>())
+    return op.getMemref().getType().getElementType();
+  llvm_unreachable("cannot find dn def");
+}
+
+template <typename T>
+static Value genConstFrom(OpBuilder &builder, Location loc, T TValue) {
   Type llvmInt32Type = builder.getIntegerType(32);
   return builder.create<LLVM::ConstantOp>(loc, llvmInt32Type,
-                                          static_cast<int32_t>(mode));
+                                          static_cast<int32_t>(TValue));
+}
+
+// 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 =
+      getCuSparseDataTypeFrom(computeTypeOptional.value_or(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();
   auto handle =
       createSparseEnvCallBuilder.create(loc, rewriter, {stream}).getResult();
   rewriter.replaceOp(op, {handle, stream});
   return success();
 }
 
 LogicalResult ConvertDestroySparseEnvOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::DestroySparseEnvOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
   auto stream = adaptor.getAsyncDependencies().front();
   Value pVec =
       MemRefDescriptor(adaptor.getMemref()).allocatedPtr(rewriter, loc);
   if (!getTypeConverter()->useOpaquePointers())
     pVec = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, pVec);
-  Type dType =
-      llvm::cast<MemRefType>(op.getMemref().getType()).getElementType();
-  auto dw = rewriter.create<LLVM::ConstantOp>(loc, llvmInt32Type,
-                                              dType.getIntOrFloatBitWidth());
+  Type dType = op.getMemref().getType().getElementType();
+  auto dtp = genConstFrom(rewriter, loc, getCuSparseDataTypeFrom(dType));
   auto handle =
       createDnVecCallBuilder
-          .create(loc, rewriter, {adaptor.getSize(), pVec, dw, stream})
+          .create(loc, rewriter, {adaptor.getSize(), pVec, dtp, stream})
           .getResult();
   rewriter.replaceOp(op, {handle, stream});
   return success();
 }
 
 LogicalResult ConvertDestroyDnVecOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::DestroyDnVecOp 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();
   Value pMat =
       MemRefDescriptor(adaptor.getMemref()).allocatedPtr(rewriter, loc);
   if (!getTypeConverter()->useOpaquePointers())
     pMat = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, pMat);
-  Type dType =
-      llvm::cast<MemRefType>(op.getMemref().getType()).getElementType();
-  auto dw = rewriter.create<LLVM::ConstantOp>(loc, llvmInt32Type,
-                                              dType.getIntOrFloatBitWidth());
+  Type dType = op.getMemref().getType().getElementType();
+  auto dtp = genConstFrom(rewriter, loc, getCuSparseDataTypeFrom(dType));
   auto handle =
       createDnMatCallBuilder
           .create(loc, rewriter,
-                  {adaptor.getRows(), adaptor.getCols(), pMat, dw, stream})
+                  {adaptor.getRows(), adaptor.getCols(), pMat, dtp, stream})
           .getResult();
   rewriter.replaceOp(op, {handle, stream});
   return success();
 }
 
 LogicalResult ConvertDestroyDnMatOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::DestroyDnMatOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
if (!getTypeConverter()->useOpaquePointers()) {
     pRowIdxs = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, pRowIdxs);
     pColIdxs = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, pColIdxs);
     pValues = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, pValues);
   }
   Type iType =
       llvm::cast<MemRefType>(op.getColIdxs().getType()).getElementType();
   Type dType =
       llvm::cast<MemRefType>(op.getValues().getType()).getElementType();
-  auto iw = rewriter.create<LLVM::ConstantOp>(
-      loc, llvmInt32Type, iType.isIndex() ? 64 : iType.getIntOrFloatBitWidth());
-  auto dw = rewriter.create<LLVM::ConstantOp>(loc, llvmInt32Type,
-                                              dType.getIntOrFloatBitWidth());
+  auto itp = genConstFrom(rewriter, loc, getCuSparseIndexTypeFrom(iType));
+  auto dtp = genConstFrom(rewriter, loc, getCuSparseDataTypeFrom(dType));
   auto handle =
       createCooCallBuilder
           .create(loc, rewriter,
                   {adaptor.getRows(), adaptor.getCols(), adaptor.getNnz(),
-                   pRowIdxs, pColIdxs, pValues, iw, dw, stream})
+                   pRowIdxs, pColIdxs, pValues, itp, dtp, stream})
           .getResult();
   rewriter.replaceOp(op, {handle, stream});
   return success();
 }
 
 LogicalResult ConvertCreateCsrOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::CreateCsrOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
if (!getTypeConverter()->useOpaquePointers()) {
     pValues = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, pValues);
   }
   Type pType =
       llvm::cast<MemRefType>(op.getRowPos().getType()).getElementType();
   Type iType =
       llvm::cast<MemRefType>(op.getColIdxs().getType()).getElementType();
   Type dType =
       llvm::cast<MemRefType>(op.getValues().getType()).getElementType();
-  auto pw = rewriter.create<LLVM::ConstantOp>(
-      loc, llvmInt32Type, pType.isIndex() ? 64 : pType.getIntOrFloatBitWidth());
-  auto iw = rewriter.create<LLVM::ConstantOp>(
-      loc, llvmInt32Type, iType.isIndex() ? 64 : iType.getIntOrFloatBitWidth());
-  auto dw = rewriter.create<LLVM::ConstantOp>(loc, llvmInt32Type,
-                                              dType.getIntOrFloatBitWidth());
+  auto ptp = genConstFrom(rewriter, loc, getCuSparseIndexTypeFrom(pType));
+  auto itp = genConstFrom(rewriter, loc, getCuSparseIndexTypeFrom(iType));
+  auto dtp = genConstFrom(rewriter, loc, getCuSparseDataTypeFrom(dType));
   auto handle =
       createCsrCallBuilder
           .create(loc, rewriter,
                   {adaptor.getRows(), adaptor.getCols(), adaptor.getNnz(),
-                   pRowPos, pColIdxs, pValues, pw, iw, dw, stream})
+                   pRowPos, pColIdxs, pValues, ptp, itp, dtp, stream})
           .getResult();
   rewriter.replaceOp(op, {handle, stream});
   return success();
 }
 
 LogicalResult ConvertDestroySpMatOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::DestroySpMatOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {

 LogicalResult ConvertSpMVBufferSizeOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::SpMVBufferSizeOp 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, 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(op.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(), 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(op.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,
+                          adaptor.getDnX(), adaptor.getDnY(), 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(op.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(), 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(op.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(), 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(op.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(), 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(op.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(),
+                           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

 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/ExecutionEngine/CRunnerUtils.h"
 
 #include <stdio.h>
 
 #include "cuda.h"
+#include "cuda_bf16.h"
+#include "cuda_fp16.h"
 #include "cusparse.h"
 
 #ifdef _WIN32
 #define MLIR_CUDA_WRAPPERS_EXPORT __declspec(dllexport)
 #else
 #define MLIR_CUDA_WRAPPERS_EXPORT
 #endif // _WIN32
 

 
 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;
-  }
-}
-
 // Some macro magic to get float/double alpha and beta on host.
-#define ALPHABETA(w, alpha, beta)                                              \
+#define ALPHABETA(dtp, alpha, beta)                                            \
+  __nv_bfloat16(alpha##bf16) = 1.0f;                                           \
+  __nv_bfloat16(beta##bf16) = 1.0f;                                            \
+  __half(alpha##f16) = 1.0f;                                                   \
+  __half(beta##f16) = 1.0f;                                                    \
   float(alpha##f) = 1.0f;                                                      \
   float(beta##f) = 1.0f;                                                       \
   double(alpha##d) = 1.0;                                                      \
   double(beta##d) = 1.0;                                                       \
   const void *(alpha##p) = nullptr;                                            \
   const void *(beta##p) = nullptr;                                             \
-  if ((w) == 32) {                                                             \
+  if (dtp == CUDA_R_16BF || dtp == CUDA_C_16BF) {                              \
+    (alpha##p) = reinterpret_cast<void *>(&(alpha##16bf));                     \
+    (beta##p) = reinterpret_cast<void *>(&(beta##16bf));                       \
+  } else if (dtp == CUDA_R_16F || dtp == CUDA_C_16F) {                         \
+    (alpha##p) = reinterpret_cast<void *>(&(alpha##16f));                      \
+    (beta##p) = reinterpret_cast<void *>(&(beta##16f));                        \
+  } else if (dtp == CUDA_R_32F || dtp == CUDA_C_32F) {                         \
     (alpha##p) = reinterpret_cast<void *>(&(alpha##f));                        \
     (beta##p) = reinterpret_cast<void *>(&(beta##f));                          \
-  } else {                                                                     \
+  } else if (dtp == CUDA_R_64F || dtp == CUDA_C_64F) {                         \
     (alpha##p) = reinterpret_cast<void *>(&(alpha##d));                        \
     (beta##p) = reinterpret_cast<void *>(&(beta##d));                          \
+  } else {                                                                     \
+    llvm_unreachable("Unsupported data type");                                 \
   }
 
 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
 mgpuDestroySparseEnv(void *h, CUstream /*stream*/) {
   cusparseHandle_t handle = reinterpret_cast<cusparseHandle_t>(h);
   CUSPARSE_REPORT_IF_ERROR(cusparseDestroy(handle))
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void *
-mgpuCreateDnVec(intptr_t size, void *values, int32_t dw, CUstream /*stream*/) {
+mgpuCreateDnVec(intptr_t size, void *values, int32_t dtp, CUstream /*stream*/) {
   cusparseDnVecDescr_t vec = nullptr;
-  cudaDataType_t dtp = dataTp(dw);
-  CUSPARSE_REPORT_IF_ERROR(cusparseCreateDnVec(&vec, size, values, dtp))
+  auto dTp = static_cast<cudaDataType_t>(dtp);
+  CUSPARSE_REPORT_IF_ERROR(cusparseCreateDnVec(&vec, size, values, dTp))
   return reinterpret_cast<void *>(vec);
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
 mgpuDestroyDnVec(void *v, CUstream /*stream*/) {
   cusparseDnVecDescr_t vec = reinterpret_cast<cusparseDnVecDescr_t>(v);
   CUSPARSE_REPORT_IF_ERROR(cusparseDestroyDnVec(vec))
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void *
-mgpuCreateDnMat(intptr_t rows, intptr_t cols, void *values, int32_t dw,
+mgpuCreateDnMat(intptr_t rows, intptr_t cols, void *values, int32_t dtp,
                 CUstream /*stream*/) {
   cusparseDnMatDescr_t mat = nullptr;
-  cudaDataType_t dtp = dataTp(dw);
+  auto dTp = static_cast<cudaDataType_t>(dtp);
   CUSPARSE_REPORT_IF_ERROR(cusparseCreateDnMat(&mat, rows, cols, /*ld=*/cols,
-                                               values, dtp, CUSPARSE_ORDER_ROW))
+                                               values, dTp, CUSPARSE_ORDER_ROW))
   return reinterpret_cast<void *>(mat);
 }
 
 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))
 }
 
+// TODO: pass in int type and data type instead of width
 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,
+              void *colIdxs, void *values, int32_t itp, int32_t dtp,
               CUstream /*stream*/) {
   cusparseSpMatDescr_t mat = nullptr;
-  cusparseIndexType_t itp = idxTp(iw);
-  cudaDataType_t dtp = dataTp(dw);
+  auto iTp = static_cast<cusparseIndexType_t>(itp);
+  auto dTp = static_cast<cudaDataType_t>(dtp);
   CUSPARSE_REPORT_IF_ERROR(cusparseCreateCoo(&mat, rows, cols, nnz, rowIdxs,
-                                             colIdxs, values, itp,
-                                             CUSPARSE_INDEX_BASE_ZERO, dtp))
+                                             colIdxs, values, iTp,
+                                             CUSPARSE_INDEX_BASE_ZERO, dTp))
   return reinterpret_cast<void *>(mat);
 }
 
+// TODO: pass in int type and data type instead of width
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void *
 mgpuCreateCsr(intptr_t rows, intptr_t cols, intptr_t nnz, void *rowPos,
-              void *colIdxs, void *values, int32_t pw, int32_t iw, int32_t dw,
-              CUstream /*stream*/) {
+              void *colIdxs, void *values, int32_t ptp, int32_t itp,
+              int32_t dtp, CUstream /*stream*/) {
   cusparseSpMatDescr_t mat = nullptr;
-  cusparseIndexType_t ptp = idxTp(pw);
-  cusparseIndexType_t itp = idxTp(iw);
-  cudaDataType_t dtp = dataTp(dw);
+  auto pTp = static_cast<cusparseIndexType_t>(ptp);
+  auto iTp = static_cast<cusparseIndexType_t>(itp);
   CUSPARSE_REPORT_IF_ERROR(cusparseCreateCsr(&mat, rows, cols, nnz, rowPos,
-                                             colIdxs, values, ptp, itp,
+                                             colIdxs, values, pTp, iTp,
                                              CUSPARSE_INDEX_BASE_ZERO, dtp))
   return reinterpret_cast<void *>(mat);
 }
 
 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,
+mgpuSpMVBufferSize(void *h, int32_t ma, void *a, void *x, void *y, int32_t ctp,
                    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);
-  ALPHABETA(dw, alpha, beta)
+  cudaDataType_t cTp = static_cast<cudaDataType_t>(ctp);
+  ALPHABETA(cTp, 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))
+                              cTp, 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,
+                                                   void *x, void *y,
+                                                   int32_t ctp, 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);
-  ALPHABETA(dw, alpha, beta)
+  cudaDataType_t cTp = static_cast<cudaDataType_t>(ctp);
+  ALPHABETA(cTp, alpha, beta)
   CUSPARSE_REPORT_IF_ERROR(cusparseSpMV(handle, modeA, alphap, matA, vecX,
-                                        betap, vecY, dtp,
+                                        betap, vecY, cTp,
                                         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 ctp, 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);
-  ALPHABETA(dw, alpha, beta)
+  cudaDataType_t cTp = static_cast<cudaDataType_t>(ctp);
+  ALPHABETA(cTp, 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, cTp,
       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*/) {
+mgpuSpMM(void *h, int32_t ma, int32_t mb, void *a, void *b, void *c,
+         int32_t ctp, 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);
-  ALPHABETA(dw, alpha, beta)
+  cudaDataType_t cTp = static_cast<cudaDataType_t>(ctp);
+  ALPHABETA(cTp, alpha, beta)
   CUSPARSE_REPORT_IF_ERROR(cusparseSpMM(handle, modeA, modeB, alphap, matA,
-                                        matB, betap, matC, dtp,
+                                        matB, betap, matC, cTp,
                                         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*/) {
+                    int32_t ctp, 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);
   cusparseDnMatDescr_t matA = reinterpret_cast<cusparseDnMatDescr_t>(a);
   cusparseDnMatDescr_t matB = reinterpret_cast<cusparseDnMatDescr_t>(b);
   cusparseSpMatDescr_t matC = reinterpret_cast<cusparseSpMatDescr_t>(c);
-  cudaDataType_t dtp = dataTp(dw);
-  ALPHABETA(dw, alpha, beta)
+  auto cTp = static_cast<cudaDataType_t>(ctp);
+  ALPHABETA(cTp, alpha, beta)
   size_t bufferSize = 0;
   CUSPARSE_REPORT_IF_ERROR(cusparseSDDMM_bufferSize(
-      handle, modeA, modeB, alphap, matA, matB, betap, matC, dtp,
+      handle, modeA, modeB, alphap, matA, matB, betap, matC, cTp,
       CUSPARSE_SDDMM_ALG_DEFAULT, &bufferSize))
   return bufferSize == 0 ? 1 : bufferSize; // avoid zero-alloc
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
 mgpuSDDMM(void *h, int32_t ma, int32_t mb, void *a, void *b, void *c,
-          int32_t dw, void *buf, CUstream /*stream*/) {
+          int32_t ctp, 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);
   cusparseDnMatDescr_t matA = reinterpret_cast<cusparseDnMatDescr_t>(a);
   cusparseDnMatDescr_t matB = reinterpret_cast<cusparseDnMatDescr_t>(b);
   cusparseSpMatDescr_t matC = reinterpret_cast<cusparseSpMatDescr_t>(c);
-  cudaDataType_t dtp = dataTp(dw);
-  ALPHABETA(dw, alpha, beta)
+  auto cTp = static_cast<cudaDataType_t>(ctp);
+  ALPHABETA(cTp, alpha, beta)
   CUSPARSE_REPORT_IF_ERROR(cusparseSDDMM(handle, modeA, modeB, alphap, matA,
-                                         matB, betap, matC, dtp,
+                                         matB, betap, matC, cTp,
                                          CUSPARSE_SDDMM_ALG_DEFAULT, buf))
 }

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

+// RUN: mlir-opt %s -split-input-file | mlir-opt -split-input-file | FileCheck %s
+
+module attributes {gpu.container_module} {
+
+  // CHECK-LABEL: func @matvec
+  // CHECK: %{{.*}} = gpu.wait async
+  // CHECK: %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xindex>
+  // CHECK: %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xf64>
+  // CHECK: %{{.*}}, %{{.*}} = gpu.create_sparse_env async [%{{.*}}]
+  // CHECK: %{{.*}}, %{{.*}} = gpu.create_coo async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xindex>, memref<?xindex>, memref<?xf64>
+  // CHECK: %{{.*}}, %{{.*}} = gpu.create_dn_vec async [%{{.*}}] %{{.*}}, %{{.*}} : memref<?xf64>
+  // CHECK: %{{.*}}, %{{.*}} = gpu.spmv_buffer_size async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}
+  // CHECK: %{{.*}} = gpu.spmv async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xf64>
+  // CHECK: %{{.*}} = gpu.destroy_sp_mat async [%{{.*}}] %{{.*}}
+  // CHECK: %{{.*}} = gpu.destroy_dn_vec async [%{{.*}}] %{{.*}}
+  // CHECK: %{{.*}} = gpu.destroy_sparse_env async [%{{.*}}] %{{.*}}
+  // CHECK: gpu.wait [%{{.*}}]
+  // CHECK: return
+  func.func @matvec(%arg0: index) {
+    %token0 = gpu.wait async
+    %mem1, %token1 = gpu.alloc async [%token0] (%arg0) : memref<?xindex>
+    %mem2, %token2 = gpu.alloc async [%token1] (%arg0) : memref<?xf64>
+    %env, %token3 = gpu.create_sparse_env async [%token2]
+    %spmat, %token4 = gpu.create_coo async [%token3] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
+    %dnvec, %token5 = gpu.create_dn_vec async [%token4] %mem2, %arg0 : memref<?xf64>
+    %bufferSz, %token6 = gpu.spmv_buffer_size async [%token5] %env, %spmat, %dnvec, %dnvec
+    %token7 = gpu.spmv async [%token6] %env, %spmat, %dnvec, %dnvec, %mem2 : memref<?xf64>
+    %token8 = gpu.destroy_sp_mat async [%token7] %spmat
+    %token9 = gpu.destroy_dn_vec async [%token8] %dnvec
+    %token10 = gpu.destroy_sparse_env async [%token9] %env
+    gpu.wait [%token10]
+    return
+  }
+
+  // CHECK-LABEL: func @matmul
+  // CHECK: %{{.*}} = gpu.wait async
+  // CHECK: %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xindex>
+  // CHECK: %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xf64>
+  // CHECK: %{{.*}}, %{{.*}} = gpu.create_sparse_env async [%{{.*}}]
+  // CHECK: %{{.*}}, %{{.*}} = gpu.create_csr async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xindex>, memref<?xindex>, memref<?xf64>
+  // CHECK: %{{.*}}, %{{.*}} = gpu.create_dn_mat async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}} : memref<?xf64>
+  // CHECK: %{{.*}}, %{{.*}} = gpu.spmm_buffer_size async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : f64
+  // CHECK: %{{.*}} = gpu.spmm async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xf64>, f64
+  // CHECK: %{{.*}} = gpu.destroy_sp_mat async [%{{.*}}] %{{.*}}
+  // CHECK: %{{.*}} = gpu.destroy_dn_mat async [%{{.*}}] %{{.*}}
+  // CHECK: %{{.*}} = gpu.destroy_sparse_env async [%{{.*}}] %{{.*}}
+  // CHECK: gpu.wait [%{{.*}}]
+  // CHECK: return
+  func.func @matmul(%arg0: index) {
+    %token0 = gpu.wait async
+    %mem1, %token1 = gpu.alloc async [%token0] (%arg0) : memref<?xindex>
+    %mem2, %token2 = gpu.alloc async [%token1] (%arg0) : memref<?xf64>
+    %env, %token3 = gpu.create_sparse_env async [%token2]
+    %spmat, %token4 = gpu.create_csr async [%token3] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
+    %dnmat, %token5 = gpu.create_dn_mat async [%token4] %arg0, %arg0, %mem2 : memref<?xf64>
+    %bufferSz, %token6 = gpu.spmm_buffer_size async [%token5] %env, %spmat, %dnmat, %dnmat : f64
+    %token7 = gpu.spmm async [%token6] %env, %spmat, %dnmat, %dnmat, %mem2 : memref<?xf64>, f64
+    %token8 = gpu.destroy_sp_mat async [%token7] %spmat
+    %token9 = gpu.destroy_dn_mat async [%token8] %dnmat
+    %token10 = gpu.destroy_sparse_env async [%token9] %env
+    gpu.wait [%token10]
+    return
+  }
+
+  // CHECK-LABEL: func @sddmm
+  // CHECK: %{{.*}} = gpu.wait async
+  // CHECK: %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xindex>
+  // CHECK: %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xf64>
+  // CHECK: %{{.*}}, %{{.*}} = gpu.create_sparse_env async [%{{.*}}]
+  // CHECK: %{{.*}}, %{{.*}} = gpu.create_csr async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xindex>, memref<?xindex>, memref<?xf64>
+  // CHECK: %{{.*}}, %{{.*}} = gpu.create_dn_mat async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}} : memref<?xf64>
+  // CHECK: %{{.*}}, %{{.*}} = gpu.sddmm_buffer_size async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}
+  // CHECK: %{{.*}} = gpu.sddmm async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xf64>
+  // CHECK: %{{.*}} = gpu.destroy_sp_mat async [%{{.*}}] %{{.*}}
+  // CHECK: %{{.*}} = gpu.destroy_dn_mat async [%{{.*}}] %{{.*}}
+  // CHECK: %{{.*}} = gpu.destroy_sparse_env async [%{{.*}}] %{{.*}}
+  // CHECK: gpu.wait [%{{.*}}]
+  // CHECK: return
+  func.func @sddmm(%arg0: index) {
+    %token0 = gpu.wait async
+    %mem1, %token1 = gpu.alloc async [%token0] (%arg0) : memref<?xindex>
+    %mem2, %token2 = gpu.alloc async [%token1] (%arg0) : memref<?xf64>
+    %env, %token3 = gpu.create_sparse_env async [%token2]
+    %spmat, %token4 = gpu.create_csr async [%token3] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
+    %dnmat, %token5 = gpu.create_dn_mat async [%token4] %arg0, %arg0, %mem2 : memref<?xf64>
+    %bufferSz, %token6 = gpu.sddmm_buffer_size async [%token5] %env, %dnmat, %dnmat, %spmat
+    %token7 = gpu.sddmm async [%token6] %env, %dnmat, %dnmat, %spmat, %mem2 : memref<?xf64>
+    %token8 = gpu.destroy_sp_mat async [%token7] %spmat
+    %token9 = gpu.destroy_dn_mat async [%token8] %dnmat
+    %token10 = gpu.destroy_sparse_env async [%token9] %env
+    gpu.wait [%token10]
+    return
+  }
+
+}
+
+