[mlir][sparse][gpu] add sm8.0+ tensor core 2:4 sparsity support

[K-Wu]

seemingly finished the cuda runtime wrapper part

[K-Wu]

minor

[K-Wu]

in progress

[K-Wu]

in-progress

[K-Wu]

fix compile error ?

[K-Wu]

add tests and finished lower pass

[K-Wu]

correct type in the test

[K-Wu]

stuck at declaring tuple literal in .mlir

mlir/test/Conversion/GPUCommon/2To4Sparsity/lower-2to4-sparse-to-gpu-runtime-calls.mlir

28

How can I create a three-element tuple literal here?

[K-Wu]

correcting the type

[K-Wu]

the lowering pass now works but I added the env handle to the CreateDnMatOp

[K-Wu]

fixing compile error

[K-Wu]

fixing test script

[K-Wu]

refine names

[K-Wu]

rebase origin/main

[K-Wu]

fixing compile errors caused by rebase

[K-Wu]

fixing merging errors and remove tuple types

[K-Wu]

fixing compile errors; improve documentation

[K-Wu]

fixing merging error

[K-Wu]

fix example a bit

[K-Wu]

fixing compile errors

[K-Wu]

fixing compile errors

[K-Wu]

fix test error

[K-Wu]

clean up unnecessary test flag addition

[K-Wu]

minor

[K-Wu]

using cmake find to switch on/off cuda official libraries

[K-Wu]

mlir/lib/ExecutionEngine/CMakeLists.txt

194–195

cusparse guarded without adding new build options

216

cusparse guarded without adding new build options

[K-Wu]

fix cmakelists.txt

[K-Wu]

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

547

TODO: transpose mode A and B are specified here

549

TODO: pass in compute_type here
CUSPARSE_COMPUTE {_16F, _32I, _TF32, TF32_FAST}

[K-Wu]

rebase origin/main

[K-Wu]

implementing transpose and compute type

[K-Wu]

done type enum

[K-Wu]

format

[K-Wu]

add todo

[K-Wu]

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

552

done

mlir/test/Conversion/GPUCommon/2To4Sparsity/lower-2to4-sparse-to-gpu-runtime-calls.mlir

28

This is now Variadic and no tuple anymore as either operand or result in the newly added/changed ops

[K-Wu]

fix format

[aartbik]

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

1675

why did you add the handle to DnMat?

2006

this change seems unrelated to the core change of this revision?

mlir/lib/ExecutionEngine/CMakeLists.txt

193–194

Since you change the logic of the make I would move these two "find_lib" below to L217
Also, no longer say

"we need" as we did for L190

but now say,

"Determine if cuSPARSE lib and cuSPARSELT are available and, if so, set the MLIR_CUDA_SM80_SPARSE_ENABLED flag"

216

so right before this, the two find_lib without REQUIRED now

233

We now have situations like

(1) we cannot find either lib
(2) we find cusparse but not cusparselt
(3) we find both

do we simply want to ONLY support (3) and guard all code with the MLIR_CUDA_SM80_SPARSE_ENABLED
or recognize (1) and (2) with a MLIR_CUDA_SPARSE_ENABLED

?

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

24–29

since we now have builds possible without cusparse lib, don't we need a macro for this (and all cusparse wrappers?

448

I would move this to L438, and use two empty lines around the #if since it applies to rest of file

593

newline?

[K-Wu]

mlir/lib/ExecutionEngine/CMakeLists.txt

233

Working on it. Adding a second flag in the wrapper cpp where CUSPARSE disabled will imply CUSPARSELT disabled

[K-Wu]

fix one of the ptr to local scope bug and addressing some comments

[K-Wu]

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

448

addressed. Let me know if your thought is different from my new changes :D

[K-Wu]

adding unsynced cmake change; addressing comments

[K-Wu]

mlir/lib/ExecutionEngine/CMakeLists.txt

193–194

Sorry I failed to sync changes from another machine which I thought had been synced by arcanist.

Please check the following few lines for the improved CMake build flag logic.

Please also check if the runtime flags in the CUDARuntimeWrappers.cpp address the comments.

[K-Wu]

recover unsynced changes

[K-Wu]

impl alloca for handles

[K-Wu]

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

1675

cusparseLt's API is a bit different from cusparse and requires environment handle as input.

But I am open to ideas to split the revision. Also it seems feasible to retrieve the env implicitly as well. Let's discuss offline maybe.

2006

When using cusparseLt I get three sizes when invoking spmm_buffer_size. THe three sizes are for workspace, compressed data and compressed buffer. Tablegen asks for a type in the assembly to understand whether the output is one size, or a tuple of three sizes. Let us discuss offline how to improve this.

[K-Wu]

rebase origin/main

[K-Wu]

fix unsynced test error. i hate arcanist

[K-Wu]

format

[K-Wu]

fix typo

[aartbik]

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

1675

Ah for the 2:4 case.

shall we add it to dn_vec too then, for consistency?

mlir/lib/ExecutionEngine/CMakeLists.txt

194–195

I gave this some more thought, and I don't think we have any precedent in LLVM/MLIR building conditionally on the context. In fact, some users may not like this, since exactly the same config will result in different binaries on different installations, causing perhaps hard-to-track bugs on bots.

So let's keep the two flags, but keep the build logic

if(MLIR_ENABLE_CUDA_CUSPARSE)

find_library(CUDA_CUSPARSE_LIBRARY cusparse HINTS ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES} REQUIRED)
 ...
 if(MLIR_ENABLE_CUDA_CUSPARSELT)
   find_library(CUDA_CUSPARSELT_LIBRARY cusparseLt HINTS ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES}  REQUIRED )
   find_path(CUDA_CUSPARSELT_HEADER cusparseLt.h HINTS ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}   REQUIRED)
   ...

Note
(1) the REQUIRED now that the build is requested by config
(2) CUSPARSELT is nested in CUSPARSE, i.e. you cannot have the LT without the non-LT version (seems rasonable)

[K-Wu]

mlir/lib/ExecutionEngine/CMakeLists.txt

194–195

That makes sense. Yeah cusparseLt uses some data structures in cusparse so (2) is very reasonable

[K-Wu]

addressing comments

[K-Wu]

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

1675

That makes sense. Done.

[K-Wu]

fixing compile error

[K-Wu]

fixing compile error

[K-Wu]

fixing compile error

[aartbik]

mlir/lib/ExecutionEngine/CMakeLists.txt

194–195

This must now go under the CUDA_CUSPARSE_LIBRARY below

otherwise, this line will fail without lib even when macro is not set

197

Try to keep the comments under 80-col by breaking the line

typo: installed

Also, rephrase. Enabling the MLIR_ENABLE_CUDA_CUSPARSELT flag assumes the library is installed ...

198

you don't want to nest this in the one below?

[K-Wu]

addressing comments

[K-Wu]

addressing comments

[K-Wu]

adding necesary bazel changes. credit to Aart

[aartbik]

mlir/lib/ExecutionEngine/CMakeLists.txt

216

This comment is outdated, since we no longer set the flag

1. Find the libcusparse.so library if CUSPARSE build is requested.

224

Find the libcusparseLt.so library if CUSPARSELT build is requested.

<then rest of your comment, but broken up with 80-cols>

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

27

I would nest this done inside CUSPARSE #if/#endif
since we only support both set, or former set only

444

move this end to end of file, so we have nested definitions

mlir/test/Conversion/GPUCommon/2To4Sparsity/lit.local.cfg

2

not needed, we are not running this, just lowering

mlir/test/Conversion/GPUCommon/2To4Sparsity/lower-2to4-sparse-to-gpu-runtime-calls.mlir

2

don't introduce a new 2To4Sparsity directory, simply move this test file up to the same level as e.g. lower-sparse-to-gpu-runtime-calls.mlir

[K-Wu]

fixing bazel errors and addressing comments

[K-Wu]

rebase origin/main

[K-Wu]

fix compile error; add moved new file

[aartbik]

did you forget to add the lowering file to this revision after moving up?

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

818

Also add CreataCooAoS for completeness?
It is deprecated on the lib side, but it still lives as op here

822

Start comments with upper case and end with period

// Print the SpMat defining op.

Here and below a few times

827

a bit strange to return this as string?
Just a bool?

[K-Wu]

fixing compile error

[K-Wu]

addressing comments

[K-Wu]

Addressed comments

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

let description = [{
 
     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.
 
     Example:
 
     ```mlir
-    %dvec, %token = gpu.create_dn_vec async [%dep] %mem, %size : memref<?xf64>
+    %dvec, %token = gpu.create_dn_vec async [%dep] %env, %mem, %size : memref<?xf64>
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
-                   AnyMemRef:$memref, Index:$size);
+                       GPU_SparseEnvHandle:$env,
+                       AnyMemRef:$memref, 
+                       Index:$size);
   let results = (outs Res<GPU_SparseDnVecHandle>:$dvec, 
                       Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $memref `,` $size attr-dict `:` type($memref)
+    $env `,` $memref `,` $size attr-dict `:` type($memref)
   }];
 }
 
 def GPU_DestroyDnVecOp : GPU_Op<"destroy_dn_vec", [GPU_AsyncOpInterface]> {
   let summary = "Destroy dense vector operation";
   let description = [{
     The `gpu.destroy_dn_vec` operation releases all resources of a dense
     vector represented by a handle that was previously created by a
let description = [{
 
     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.
 
     Example:
 
     ```mlir
-    %dmat, %token = gpu.create_dn_mat async [%dep] %mem, %size : memref<?xf64>
+    %dmat, %token = gpu.create_dn_mat async [%dep] %env, %rows, %cols, %mem : memref<?xf64>
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
+                       GPU_SparseEnvHandle:$env,
                        Index:$rows,
                        Index:$cols,
                        AnyMemRef:$memref);
   let results = (outs Res<GPU_SparseDnMatHandle>:$dmat, Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $rows `,` $cols `,` $memref attr-dict `:` type($memref)
+    $env `,` $rows `,` $cols `,` $memref attr-dict `:` type($memref)
   }];
 }
 
 def GPU_DestroyDnMatOp : GPU_Op<"destroy_dn_mat", [GPU_AsyncOpInterface]> {
   let summary = "Destroy dense matrix operation";
   let description = [{
     The `gpu.destroy_dn_mat` operation releases all resources of a dense
     matrix represented by a handle that was previously created by a
def GPU_CreateCsrOp : GPU_Op<"create_csr", [GPU_AsyncOpInterface]> {
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $rows `,` $cols `,` $nnz `,` $rowPos `,` $colIdxs `,` $values attr-dict
     `:` type($rowPos) `,` type($colIdxs) `,` type($values)
   }];
 }
 
+
+def GPU_Create2To4SpMatOp : GPU_Op<"create_2to4_spmat", [GPU_AsyncOpInterface]> {
+  let summary = "Create sparse matrix with 2:4 sparsity operation";
+  let description = [{
+    The `gpu.create_2to4_spmat` operation initializes a sparse matrix in dense 
+    format with 2:4 sparsity.
+    The buffers must already be copied from the host to the device prior to
+    using this operation. The operation returns a handle to the sparse
+    matrix descriptor.
+
+    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.
+
+    Example:
+
+    ```mlir
+    %spmat, %token = gpu.create_2to4_spmat async [%dep] %env, %rows, %cols, %mem : memref<?xf64>
+    ```
+  }];
+
+  let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
+                       GPU_SparseEnvHandle:$env,
+                       Index:$rows,
+                       Index:$cols,
+                       AnyMemRef:$memref);
+  let results = (outs Res<GPU_SparseSpMatHandle>:$spMat, 
+                      Optional<GPU_AsyncToken>:$asyncToken);
+
+  let assemblyFormat = [{
+    custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
+    $env `,` $rows `,` $cols `,` $memref attr-dict `:` type($memref)
+  }];
+}
+
 def GPU_DestroySpMatOp : GPU_Op<"destroy_sp_mat", [GPU_AsyncOpInterface]> {
   let summary = "Destroy sparse matrix operation";
   let description = [{
     The `gpu.destroy_sp_mat` operation releases all resources of a sparse
     matrix represented by a handle that was previously created by a
     one of the sparse matrix creation operations.
 
     If the `async` keyword is present, the op is executed asynchronously (i.e.
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.spmm_buffer_size async [%dep] %env, %spmatA{TRANSPOSE}, %dnmatB{TRANSPOSE}, %dnmatC into f32
+    %bufferszs, %token = gpu.spmm_buffer_size async [%dep] %env, %spmatA{TRANSPOSE}, %dnmatB{TRANSPOSE}, %dnmatC : i64 into f32
     ```
   }];
 
   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,
                        TypeAttr:$computeType);
-  let results = (outs Res<Index>:$bufferSz, 
+  let results = (outs Res<AnyTypeOf<[Index, TupleOf<[Index, Index, 
+                                                     Index]>]>>:$bufferSzs, 
                       Optional<GPU_AsyncToken>:$asyncToken);
 
   let builders = [OpBuilder<(ins
-      "Type":$bufferSz,
+      "Type":$bufferSzs,
       "Type":$asyncToken,
       "ValueRange":$asyncDependencies,
       "Value":$env,
       "Value":$spmatA,
       "Value":$dnmatB,
       "Value":$dnmatC,
       "Type":$computeType), [{
     auto modeA = gpu::TransposeMode::NON_TRANSPOSE;
     auto modeB = gpu::TransposeMode::NON_TRANSPOSE;
-    return build($_builder, $_state, bufferSz, asyncToken, asyncDependencies, 
+    return build($_builder, $_state, bufferSzs, asyncToken, asyncDependencies, 
                  env, modeA, modeB, spmatA, dnmatB, dnmatC, computeType);}]>
   ];
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $dnmatC attr-dict `into` $computeType
+    $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $dnmatC attr-dict `:` type($bufferSzs) `into` $computeType
   }];
 }
 
-def GPU_SpMMOp : GPU_Op<"spmm", [GPU_AsyncOpInterface]> {
+def GPU_SpMMOp : GPU_Op<"spmm", [GPU_AsyncOpInterface, AttrSizedOperandSegments]> {
   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
     sparse operations to construct an environment and the operands for SpMM. The
     buffer must have been allocated on the device.
 
     If the `async` keyword is present, the op is executed asynchronously (i.e.
def GPU_SpMMOp : GPU_Op<"spmm", [GPU_AsyncOpInterface, AttrSizedOperandSegments]> {
   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,
                        TypeAttr:$computeType,
-                       AnyMemRef:$buffer);
+                       Variadic<AnyMemRef>:$buffers);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let builders = [OpBuilder<(ins
       "Type":$asyncToken,
       "ValueRange":$asyncDependencies,
       "Value":$env,
       "Value":$spmatA,
       "Value":$dnmatB,
       "Value":$dnmatC,
       "Type":$computeType,
-      "Value":$buffer), [{
+      "ValueRange":$buffers), [{
     auto modeA = gpu::TransposeMode::NON_TRANSPOSE;
     auto modeB = gpu::TransposeMode::NON_TRANSPOSE;
     return build($_builder, $_state, asyncToken, asyncDependencies, env, modeA, 
-                 modeB, spmatA, dnmatB, dnmatC, computeType, buffer);}]>
+                 modeB, spmatA, dnmatB, dnmatC, computeType, buffers);}]>
   ];
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $dnmatC `,` $buffer attr-dict `:` type($buffer) `into` $computeType
+    $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $dnmatC `,` $buffers attr-dict `:` type($buffers) `into` $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.

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

protected:
   Type llvmVoidType = LLVM::LLVMVoidType::get(context);
   LLVM::LLVMPointerType llvmPointerType =
       this->getTypeConverter()->getPointerType(IntegerType::get(context, 8));
   Type llvmPointerPointerType =
       this->getTypeConverter()->getPointerType(llvmPointerType);
   Type llvmInt8Type = IntegerType::get(context, 8);
   Type llvmInt32Type = IntegerType::get(context, 32);
   Type llvmInt64Type = IntegerType::get(context, 64);
+  Type llvmInt8PointerType =
+      this->getTypeConverter()->getPointerType(llvmInt8Type);
+  Type llvmInt64PointerType =
+      this->getTypeConverter()->getPointerType(llvmInt64Type);
   Type llvmIntPtrType = IntegerType::get(
       context, this->getTypeConverter()->getPointerBitwidth(0));
 
   FunctionCallBuilder moduleLoadCallBuilder = {
       "mgpuModuleLoad",
       llvmPointerType /* void *module */,
       {llvmPointerType /* void *cubin */}};
   FunctionCallBuilder moduleUnloadCallBuilder = {
FunctionCallBuilder SDDMMBufferSizeCallBuilder = {
        llvmPointerType, llvmPointerType, llvmInt32Type,
        llvmPointerType /* void *stream */}};
   FunctionCallBuilder SDDMMCallBuilder = {
       "mgpuSDDMM",
       llvmVoidType,
       {llvmPointerType, llvmInt32Type, llvmInt32Type, llvmPointerType,
        llvmPointerType, llvmPointerType, llvmInt32Type, llvmPointerType,
        llvmPointerType /* void *stream */}};
+  FunctionCallBuilder AssertSparseLTEnvHandleSizeCallBuilder = {
+      "mgpuAssertSparseLTEnvHandleSize", llvmVoidType, {}};
+  FunctionCallBuilder AssertSparseLTSpMatHandleSizeCallBuilder = {
+      "mgpuAssertSparseLTSpMatHandleSize", llvmVoidType, {}};
+  FunctionCallBuilder AssertSparseLTDnMatHandleSizeCallBuilder = {
+      "mgpuAssertSparseLtDnMatHandleSize", llvmVoidType, {}};
+  FunctionCallBuilder createSparseLtEnvCallBuilder = {
+      "mgpuCreateSparseLtEnv",
+      llvmVoidType,
+      {llvmPointerType, llvmPointerType /* void *stream */}};
+  FunctionCallBuilder destroySparseLtEnvCallBuilder = {
+      "mgpuDestroySparseLtEnv",
+      llvmVoidType,
+      {llvmPointerType, llvmPointerType /* void *stream */}};
+  FunctionCallBuilder createLtDnMatCallBuilder = {
+      "mgpuCreateCuSparseLtDnMat",
+      llvmVoidType,
+      {llvmPointerType, llvmPointerType, llvmIntPtrType, llvmIntPtrType,
+       llvmPointerType, llvmInt32Type, llvmPointerType /* void *stream */}};
+  FunctionCallBuilder destroyCuSparseLtSpMatBuilder = {
+      "mgpuDestroyCuSparseLtSpMat",
+      llvmVoidType,
+      {llvmPointerType, llvmPointerType /* void *stream */}};
+  FunctionCallBuilder destroyCuSparseLtDnMatBuilder = {
+      "mgpuDestroyCuSparseLtDnMat",
+      llvmVoidType,
+      {llvmPointerType, llvmPointerType /* void *stream */}};
+  FunctionCallBuilder create2To4SpMatCallBuilder = {
+      "mgpuCusparseLtCreate2To4SpMat",
+      llvmVoidType,
+      {llvmPointerType, llvmPointerType, llvmIntPtrType, llvmIntPtrType,
+       llvmPointerType, llvmInt32Type, llvmPointerType /* void *stream */}};
+  FunctionCallBuilder cuSparseLtSpmmBufferSizeBuilder = {
+      "mgpuCuSparseLtSpmmBufferSize",
+      llvmVoidType,
+      {llvmPointerType, llvmPointerType, llvmPointerType,
+       llvmPointerType /*void *stream*/}};
+  FunctionCallBuilder cuSparseLtSpmmBuilder = {
+      "mgpuCuSparseLtSpmm",
+      llvmVoidType,
+      {llvmPointerType, llvmPointerType, llvmPointerType, llvmPointerType,
+       llvmInt32Type, llvmPointerType, llvmPointerType, 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)
ConvertCreateCsrOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter)
       : ConvertOpToGpuRuntimeCallPattern<gpu::CreateCsrOp>(typeConverter) {}
 
 private:
   LogicalResult
   matchAndRewrite(gpu::CreateCsrOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+class ConvertCreate2To4SpMatOpToGpuRuntimeCallPattern
+    : public ConvertOpToGpuRuntimeCallPattern<gpu::Create2To4SpMatOp> {
+public:
+  ConvertCreate2To4SpMatOpToGpuRuntimeCallPattern(
+      LLVMTypeConverter &typeConverter)
+      : ConvertOpToGpuRuntimeCallPattern<gpu::Create2To4SpMatOp>(
+            typeConverter) {}
+
+private:
+  LogicalResult
+  matchAndRewrite(gpu::Create2To4SpMatOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 class ConvertDestroySpMatOpToGpuRuntimeCallPattern
     : public ConvertOpToGpuRuntimeCallPattern<gpu::DestroySpMatOp> {
 public:
   ConvertDestroySpMatOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter)
       : ConvertOpToGpuRuntimeCallPattern<gpu::DestroySpMatOp>(typeConverter) {}
 
 private:
   LogicalResult

 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
 }
 
+static int32_t getCuSparseLtDataTypeFrom(Type type) {
+  if (type.isF16())
+    return 0; // CUSPARSE_COMPUTE_16F,
+  if (type.isInteger(32))
+    return 1; // CUSPARSE_COMPUTE_32I
+  llvm_unreachable("unsupported type");
+  // TODO: add support to TF32
+}
+
 // 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())
static int32_t getCuSparseDataTypeFrom(Type type) {
   if (type.isInteger(16))
     return 20; // CUDA_R_16I
   if (type.isInteger(32))
     return 10; // CUDA_R_32I
 
   llvm_unreachable("unsupported element type");
 }
 
+// TODO:  We may want a run-time (of the mlir compiler) disablement/warning:
+// cusparseLt currently won't work for cuda architecture <8.0 and will trigger a
+// runtime (of the CUDA program) error , but it might be great if we could at
+// least output a warning when we found the target architecture is <8.0 and the
+// user still wants to use cusparseLt. to make sure when lowering gpu sparse
+// dialect to llvm calls, the cusparselt calls are disabled for cuda
+// architecture <8.0
+static bool is2To4Sparsity(Value spMat) {
+  if (auto op = spMat.getDefiningOp<gpu::Create2To4SpMatOp>())
+    return true;
+  if (auto op = spMat.getDefiningOp<gpu::CreateCooOp>())
+    return false;
+  if (auto op = spMat.getDefiningOp<gpu::CreateCsrOp>())
+    return false;
+  // print the spMat defining op
+  spMat.getDefiningOp()->print(llvm::errs());
+  llvm_unreachable("cannot find spmat def");
+}
+
+static std::string inferSpMMType(Value op) {
+  for (Operation *user : op.getUsers()) {
+    auto spmmOp = dyn_cast<gpu::SpMMOp>(user);
+    // if the other operator is 50% sparsity then we should use cusparseLt
+    if (!spmmOp)
+      continue;
+    if (is2To4Sparsity(spmmOp.getSpmatA()))
+      return "cusparseLt";
+  }
+  return "cusparse";
+}
+
 // 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.");

 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 =
+  // use the cusparseLt create call if the dnmat is used with spmat with
+  // 2:4 sparsity
+  Value handle;
+  if (inferSpMMType(op.getEnv()) == "cusparseLt") {
+    // asserting the size is 11024 bytes
+    AssertSparseLTEnvHandleSizeCallBuilder.create(loc, rewriter, {});
+    auto handleSz = rewriter.create<LLVM::ConstantOp>(
+        loc, getIndexType(), rewriter.getIndexAttr(11024));
+    handle = rewriter.create<LLVM::AllocaOp>(loc, llvmInt8PointerType,
+                                             llvmInt8Type, handleSz);
+    handle = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, handle);
+    createSparseLtEnvCallBuilder.create(loc, rewriter, {handle, stream})
+        .getResult();
+  } else {
+    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();
   auto stream = adaptor.getAsyncDependencies().front();
-  destroySparseEnvCallBuilder.create(loc, rewriter, {adaptor.getEnv(), stream});
+  // use the cusparseLt destroy call if the dnmat is used with spmat with
+  // 2:4 sparsity
+  if (inferSpMMType(op.getEnv()) == "cusparseLt") {
+    destroySparseLtEnvCallBuilder.create(loc, rewriter,
+                                         {adaptor.getEnv(), stream});
+  } else {
+    destroySparseEnvCallBuilder.create(loc, rewriter,
+                                       {adaptor.getEnv(), stream});
+  }
   rewriter.replaceOp(op, {stream});
   return success();
 }
 
 LogicalResult ConvertCreateDnVecOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::CreateDnVecOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
LogicalResult ConvertCreateDnMatOpToGpuRuntimeCallPattern::matchAndRewrite(
   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 = op.getMemref().getType().getElementType();
   auto dtp = genConstInt32From(rewriter, loc, getCuSparseDataTypeFrom(dType));
-  auto handle =
+  // TODO: For now, we track the use of the handle and lower it to cusparse /
+  // cusparseLt accordingly. If in a block, both cusparse and cusparseLt are
+  // used, we require two separate Creation ops to be the correct logic. In
+  // future, we may add support to using one handle in sparse tensor / GPU
+  // dialect in both cusparse and cusparseLt. use the cusparseLt create call if
+  // the dnmat is used with spmat with 2:4 sparsity
+  Value handle;
+  if (inferSpMMType(op.getDmat()) == "cusparseLt") {
+    auto envHandle = adaptor.getEnv();
+    AssertSparseLTDnMatHandleSizeCallBuilder.create(loc, rewriter, {});
+    auto handleSz = rewriter.create<LLVM::ConstantOp>(
+        loc, getIndexType(), rewriter.getIndexAttr(11032));
+    handle = rewriter.create<LLVM::AllocaOp>(loc, llvmInt8PointerType,
+                                             llvmInt8Type, handleSz);
+    handle = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, handle);
+
+    createLtDnMatCallBuilder
+        .create(loc, rewriter,
+                {handle, envHandle, adaptor.getRows(), adaptor.getCols(), pMat,
+                 dtp, stream})
+        .getResult();
+  } else {
+    handle =
         createDnMatCallBuilder
             .create(loc, rewriter,
                     {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 (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
   auto stream = adaptor.getAsyncDependencies().front();
+  // use the cusparseLt destroy call if the dnmat is used with spmat with
+  // 2:4 sparsity
+  if (inferSpMMType(op.getDmat()) == "cusparseLt") {
+    destroyCuSparseLtDnMatBuilder.create(loc, rewriter,
+                                         {adaptor.getDmat(), stream});
+  } else {
     destroyDnMatCallBuilder.create(loc, rewriter, {adaptor.getDmat(), stream});
+  }
   rewriter.replaceOp(op, {stream});
   return success();
 }
 
 LogicalResult ConvertCreateCooOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::CreateCooOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
auto handle =
           .create(loc, rewriter,
                   {adaptor.getRows(), adaptor.getCols(), adaptor.getNnz(),
                    pRowPos, pColIdxs, pValues, ptp, itp, dtp, stream})
           .getResult();
   rewriter.replaceOp(op, {handle, stream});
   return success();
 }
 
+LogicalResult ConvertCreate2To4SpMatOpToGpuRuntimeCallPattern::matchAndRewrite(
+    gpu::Create2To4SpMatOp 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 dtp = genConstInt32From(rewriter, loc, getCuSparseLtDataTypeFrom(dType));
+  auto envHandle = adaptor.getEnv();
+
+  AssertSparseLTSpMatHandleSizeCallBuilder.create(loc, rewriter, {});
+  auto handleSz = rewriter.create<LLVM::ConstantOp>(
+      loc, getIndexType(), rewriter.getIndexAttr(44104));
+  Value handle = rewriter.create<LLVM::AllocaOp>(loc, llvmInt8PointerType,
+                                                 llvmInt8Type, handleSz);
+  handle = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, handle);
+
+  create2To4SpMatCallBuilder
+      .create(loc, rewriter,
+              {handle, envHandle, adaptor.getRows(), adaptor.getCols(), pMat,
+               dtp, stream})
+      .getResult();
+  rewriter.replaceOp(op, {handle, stream});
+  return success();
+}
+
 LogicalResult ConvertDestroySpMatOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::DestroySpMatOp 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();
+  // use the cusparseLt destroy call if the spmat is 2:4 sparsity
+  if (is2To4Sparsity(op.getSpmat())) {
+    destroyCuSparseLtSpMatBuilder.create(loc, rewriter,
+                                         {adaptor.getSpmat(), stream});
+
+  } else {
     destroySpMatCallBuilder.create(loc, rewriter, {adaptor.getSpmat(), stream});
+  }
   rewriter.replaceOp(op, {stream});
   return success();
 }
 
 LogicalResult ConvertSpMVBufferSizeOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::SpMVBufferSizeOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
   auto modeA = genConstInt32From(rewriter, loc, adaptor.getModeA());
   auto modeB = genConstInt32From(rewriter, loc, adaptor.getModeB());
   auto stream = adaptor.getAsyncDependencies().front();
   auto computeType =
       genConstInt32FromComputeMode(rewriter, loc, adaptor.getComputeType());
+  Value bufferSize;
+  if (is2To4Sparsity(op.getSpmatA())) {
+    auto three = rewriter.create<LLVM::ConstantOp>(loc, getIndexType(),
+                                                   rewriter.getIndexAttr(3));
+    bufferSize = rewriter.create<LLVM::AllocaOp>(loc, llvmInt64PointerType,
+                                                 llvmInt64Type, three);
+    bufferSize =
+        rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, bufferSize);
 
-  auto bufferSize = spMMBufferSizeCallBuilder
-                        .create(loc, rewriter,
-                                {adaptor.getEnv(), modeA, modeB,
-                                 adaptor.getSpmatA(), adaptor.getDnmatB(),
-                                 adaptor.getDnmatC(), computeType, stream})
-                        .getResult();
+    cuSparseLtSpmmBufferSizeBuilder
+        .create(loc, rewriter,
+                {bufferSize, adaptor.getEnv(), adaptor.getSpmatA(), stream})
+        .getResult();
+    rewriter.replaceOp(op, {bufferSize, stream});
+  } else {
+    bufferSize = spMMBufferSizeCallBuilder
     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)))
if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
     return failure();
   Location loc = op.getLoc();
   auto modeA = genConstInt32From(rewriter, loc, adaptor.getModeA());
   auto modeB = genConstInt32From(rewriter, loc, adaptor.getModeB());
   auto computeType =
       genConstInt32FromComputeMode(rewriter, loc, adaptor.getComputeType());
 
   auto stream = adaptor.getAsyncDependencies().front();
-  Value pBuf =
-      MemRefDescriptor(adaptor.getBuffer()).allocatedPtr(rewriter, loc);
+
+  // lower to cusparseLt if applicable
+  if (is2To4Sparsity(op.getSpmatA())) {
+    SmallVector<Value> pBufs;
+    for (Value buffer : adaptor.getBuffers()) {
+      Value pBuf = MemRefDescriptor(buffer).allocatedPtr(rewriter, loc);
+      if (!getTypeConverter()->useOpaquePointers())
+        pBuf = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, pBuf);
+      pBufs.push_back(pBuf);
+    }
+    cuSparseLtSpmmBuilder.create(loc, rewriter,
+                                 {adaptor.getEnv(), adaptor.getSpmatA(),
+                                  adaptor.getDnmatB(), adaptor.getDnmatC(),
+                                  computeType, pBufs[0], pBufs[1], pBufs[2],
+                                  stream});
+  } else {
+    Value pBuf = MemRefDescriptor(adaptor.getBuffers().front())
+                     .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(), computeType,
-                          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(
patterns.add<ConvertAllocOpToGpuRuntimeCallPattern,
                ConvertCreateSparseEnvOpToGpuRuntimeCallPattern,
                ConvertDestroySparseEnvOpToGpuRuntimeCallPattern,
                ConvertCreateDnVecOpToGpuRuntimeCallPattern,
                ConvertDestroyDnVecOpToGpuRuntimeCallPattern,
                ConvertCreateDnMatOpToGpuRuntimeCallPattern,
                ConvertDestroyDnMatOpToGpuRuntimeCallPattern,
                ConvertCreateCooOpToGpuRuntimeCallPattern,
                ConvertCreateCsrOpToGpuRuntimeCallPattern,
+               ConvertCreate2To4SpMatOpToGpuRuntimeCallPattern,
                ConvertDestroySpMatOpToGpuRuntimeCallPattern,
                ConvertSpMVBufferSizeOpToGpuRuntimeCallPattern,
                ConvertSpMVOpToGpuRuntimeCallPattern,
                ConvertSpMMBufferSizeOpToGpuRuntimeCallPattern,
                ConvertSpMMOpToGpuRuntimeCallPattern,
                ConvertSDDMMBufferSizeOpToGpuRuntimeCallPattern,
                ConvertSDDMMOpToGpuRuntimeCallPattern>(converter);
   patterns.add<ConvertLaunchFuncOpToGpuRuntimeCallPattern>(
       converter, gpuBinaryAnnotation, kernelBarePtrCallConv);
   patterns.add<EraseGpuModuleOpPattern>(&converter.getContext());
 }

mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp

static LogicalResult rewriteSpMV(PatternRewriter &rewriter,
   Value handle = env.getResult(0);
   token = env.getAsyncToken();
   Operation *spGenA =
       genSpMat(rewriter, loc, spmatHandleTp, tokenTp, token, szY, szX, nseA,
                rowA, colA, valA, isCOO, enableRT);
   Value spMatA = spGenA->getResult(0);
   token = spGenA->getResult(1);
   auto dvecX = rewriter.create<gpu::CreateDnVecOp>(loc, dnVecHandleTp, tokenTp,
-                                                   token, vecX, szX);
+                                                   token, handle, vecX, szX);
   Value dnX = dvecX.getResult(0);
   token = dvecX.getAsyncToken();
   auto dvecY = rewriter.create<gpu::CreateDnVecOp>(loc, dnVecHandleTp, tokenTp,
-                                                   token, vecY, szY);
+                                                   token, handle, vecY, szY);
   Value dnY = dvecY.getResult(0);
   token = dvecY.getAsyncToken();
 
   auto dnYType = llvm::cast<ShapedType>(y.getType()).getElementType();
 
   // Precompute buffersize for SpMV.
   auto bufferComp = rewriter.create<gpu::SpMVBufferSizeOp>(
       loc, indexTp, tokenTp, token, handle, spMatA, dnX, dnY,
auto env =
       rewriter.create<gpu::CreateSparseEnvOp>(loc, envHandleTp, tokenTp, token);
   Value handle = env.getResult(0);
   token = env.getAsyncToken();
   Operation *spGenA =
       genSpMat(rewriter, loc, spMatHandleTp, tokenTp, token, szm, szk, nseA,
                rowA, colA, valA, isCOO, enableRT);
   Value spMatA = spGenA->getResult(0);
   token = spGenA->getResult(1);
-  auto dmatB = rewriter.create<gpu::CreateDnMatOp>(loc, dnMatHandleTp, tokenTp,
-                                                   token, szk, szn, matB);
+  auto dmatB = rewriter.create<gpu::CreateDnMatOp>(
+      loc, dnMatHandleTp, tokenTp, token, handle, szk, szn, matB);
   Value dnB = dmatB.getResult(0);
   token = dmatB.getAsyncToken();
-  auto dmatC = rewriter.create<gpu::CreateDnMatOp>(loc, dnMatHandleTp, tokenTp,
-                                                   token, szm, szn, matC);
+  auto dmatC = rewriter.create<gpu::CreateDnMatOp>(
+      loc, dnMatHandleTp, tokenTp, token, handle, szm, szn, matC);
   Value dnC = dmatC.getResult(0);
   token = dmatC.getAsyncToken();
 
   auto dmatCType = llvm::cast<ShapedType>(c.getType()).getElementType();
 
   // Precompute buffersize for SpMM.
   auto bufferComp = rewriter.create<gpu::SpMMBufferSizeOp>(
       loc, indexTp, tokenTp, token, handle, spMatA, dnB, dnC,

mlir/lib/ExecutionEngine/CMakeLists.txt

if (CMAKE_CUDA_COMPILER)
       enable_language(CUDA)
     else()
       message(SEND_ERROR
         "Building the mlir cuda runner requires a working CUDA install")
     endif()
 
     # We need the libcuda.so library.
     find_library(CUDA_RUNTIME_LIBRARY cuda HINTS ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES} REQUIRED)
 
-    # We need the libcusparse.so library.
-    find_library(CUDA_CUSPARSE_LIBRARY cusparse HINTS ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES} REQUIRED)
+
+
 
     add_mlir_library(mlir_cuda_runtime
       SHARED
       CudaRuntimeWrappers.cpp
 
       EXCLUDE_FROM_LIBMLIR
     )
     set_property(TARGET mlir_cuda_runtime PROPERTY CXX_STANDARD 14)
+
     target_include_directories(mlir_cuda_runtime
       PRIVATE
       ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}
     )
     target_link_libraries(mlir_cuda_runtime
       PRIVATE
       ${CUDA_RUNTIME_LIBRARY}
+    )
+
+    if(MLIR_ENABLE_CUDA_CUSPARSE)
+      
+      # Determine if cuSPARSE lib and cuSPARSELT are available and, if so, set the MLIR_CUDA_SM80_SPARSE_ENABLED flag
+      find_library(CUDA_CUSPARSE_LIBRARY cusparse HINTS ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES} REQUIRED)
+      target_link_libraries(mlir_cuda_runtime
+        PRIVATE
         ${CUDA_CUSPARSE_LIBRARY}
       )
+
+      if(MLIR_ENABLE_CUDA_CUSPARSELT)
+        #libcusparseLt.so provides sm80+ tensor core support for 2:4 sparsity acceleration. Enabling the MLIR_ENABLE_CUDA_CUSPARSELT flag assumes the library is installed .
+        find_library(CUDA_CUSPARSELT_LIBRARY cusparseLt HINTS ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES} REQUIRED)
+        find_path(CUDA_CUSPARSELT_HEADER cusparseLt.h HINTS ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} REQUIRED)
+        target_include_directories(mlir_cuda_runtime
+          PRIVATE
+          ${CUDA_CUSPARSELT_HEADER}
+        )
+        target_link_libraries(mlir_cuda_runtime
+          PRIVATE
+          ${CUDA_CUSPARSELT_LIBRARY}
+        )
+      endif()
+    endif()
   endif()
 
   if(MLIR_ENABLE_ROCM_RUNNER)
     # Configure ROCm support.
     if (NOT DEFINED ROCM_PATH)
       if (NOT DEFINED ENV{ROCM_PATH})
         set(ROCM_PATH "/opt/rocm" CACHE PATH "Path to which ROCm has been installed")
       else()

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

 
 #include "mlir/ExecutionEngine/CRunnerUtils.h"
 
 #include <stdio.h>
 
 #include "cuda.h"
 #include "cuda_bf16.h"
 #include "cuda_fp16.h"
+
+#if MLIR_ENABLE_CUDA_CUSPARSE
 #include "cusparse.h"
+#endif // MLIR_ENABLE_CUDA_CUSPARSE
+
+#if MLIR_ENABLE_CUDA_CUSPARSELT
+#include "cusparseLt.h"
+#endif // MLIR_ENABLE_CUDA_CUSPARSELT
 
 #ifdef _WIN32
 #define MLIR_CUDA_WRAPPERS_EXPORT __declspec(dllexport)
 #else
 #define MLIR_CUDA_WRAPPERS_EXPORT
 #endif // _WIN32
 
 #define CUDA_REPORT_IF_ERROR(expr)                                             \
mgpuMemHostUnregisterMemRef(int64_t rank,
   auto *ptr = descriptor->data + descriptor->offset * elementSizeBytes;
   mgpuMemHostUnregister(ptr);
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSetDefaultDevice(int32_t device) {
   defaultDevice = device;
 }
 
+#if MLIR_ENABLE_CUDA_CUSPARSE
+
 ///
 /// Wrapper methods for the cuSparse library.
 ///
 
 // Some macro magic to get float/double alpha and beta on host.
 #define ALPHABETA(dtp, alpha, beta)                                            \
   __nv_bfloat16(alpha##16bf) = 1.0f;                                           \
   __nv_bfloat16(beta##16bf) = 1.0f;                                            \
mgpuSDDMM(void *h, int32_t ma, int32_t mb, void *a, void *b, void *c,
   cusparseDnMatDescr_t matB = reinterpret_cast<cusparseDnMatDescr_t>(b);
   cusparseSpMatDescr_t matC = reinterpret_cast<cusparseSpMatDescr_t>(c);
   auto cTp = static_cast<cudaDataType_t>(ctp);
   ALPHABETA(cTp, alpha, beta)
   CUSPARSE_REPORT_IF_ERROR(cusparseSDDMM(handle, modeA, modeB, alphap, matA,
                                          matB, betap, matC, cTp,
                                          CUSPARSE_SDDMM_ALG_DEFAULT, buf))
 }
+#endif // MLIR_ENABLE_CUDA_CUSPARSE
+
+#if MLIR_ENABLE_CUDA_CUSPARSELT
+
+///
+/// Wrapper methods for the cuSparseLt library.
+///
+
+struct cusparseLtSpMatHandleAndData {
+  cusparseLtMatDescriptor_t mat;
+  void *values{nullptr};
+  // TODO: the following is associated with the SpMM operator rather than the
+  // sparse matrix. Create workspace buffers and pass them to the SpMM
+  // execution.
+  cusparseLtMatmulAlgSelection_t alg_sel;
+  cusparseLtMatmulPlan_t plan;
+  cusparseLtMatmulDescriptor_t matmul;
+};
+
+struct cusparseLtDnMatHandleAndData {
+  cusparseLtMatDescriptor_t mat;
+  void *values{nullptr};
+};
+
+extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuAssertSparseLTEnvHandleSize() {
+  assert(sizeof(cusparseLtHandle_t) == 11024);
+}
+
+extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuAssertSparseLtSpMatHandleSize() {
+  return assert(sizeof(cusparseLtSpMatHandleAndData) == 44104);
+}
+
+extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSparseLtDnMatHandleSize() {
+  return assert(sizeof(cusparseLtDnMatHandleAndData) == 11032);
+}
+
+extern "C" MLIR_CUDA_WRAPPERS_EXPORT void *
+mgpuCreateSparseLtEnv(void *h, CUstream /*stream*/) {
+  // note that cuSparseLt still uses cusparseStatus_t
+  CUSPARSE_REPORT_IF_ERROR(
+      cusparseLtInit(reinterpret_cast<cusparseLtHandle_t *>(h)))
+  return;
+}
+
+extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
+mgpuDestroySparseLtEnv(void *h, CUstream /*stream*/) {
+  auto handle = reinterpret_cast<cusparseLtHandle_t *>(h);
+  CUSPARSE_REPORT_IF_ERROR(cusparseLtDestroy(handle))
+}
+
+extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
+mgpuCreateCuSparseLtDnMat(void *dh, void *h, intptr_t rows, intptr_t cols,
+                          void *values, int32_t dw, CUstream /*stream*/) {
+  cusparseLtMatDescriptor_t mat;
+  auto handle = reinterpret_cast<cusparseLtHandle_t *>(h);
+  auto dnmat_handle = reinterpret_cast<cusparseLtDnMatHandleAndData *>(dh);
+  cudaDataType_t dtp = dataTp(dw);
+  // assuming row-major when deciding lda
+  CUSPARSE_REPORT_IF_ERROR(cusparseLtDenseDescriptorInit(
+      handle, &(dh->mat), rows, cols, /*lda=*/cols,
+      /*alignment=*/16, dtp, CUSPARSE_ORDER_ROW))
+  dnmat_handle->values = values;
+}
+
+// This can be used to destroy both dense matrices and sparse matrices in
+// cusparseLt
+extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
+mgpuDestroyCuSparseLtSpMat(void *m, CUstream /*stream*/) {
+  auto matAndData = reinterpret_cast<cusparseLtSpMatHandleAndData>(m);
+}
+
+extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
+mgpuDestroyCuSparseLtDnMat(void *m, CUstream /*stream*/) {
+  auto matAndData = reinterpret_cast<cusparseLtDnMatHandleAndData>(m);
+  CUSPARSE_REPORT_IF_ERROR(cusparseLtMatDescriptorDestroy(&(mat->mat)))
+}
+
+extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
+mgpuCusparseLtCreate2To4SpMat(void *sh, void *h, intptr_t rows, intptr_t cols,
+                              void *values, int32_t dw, CUstream /*stream*/) {
+  auto spmat_handle = reinterpret_cast<cusparseLtSpMatHandleAndData *>(sh);
+  spmat_handle->values = values;
+  auto handle = reinterpret_cast<cusparseLtHandle_t *>(h);
+  cudaDataType_t dtp = dataTp_cusparseLt(dw);
+  // assuming row-major when deciding lda
+  CUSPARSE_REPORT_IF_ERROR(cusparseLtStructuredDescriptorInit(
+      handle, &(sh->mat), rows, cols, /*ld=*/cols, /*alignment=*/16, dtp,
+      CUSPARSE_ORDER_ROW, CUSPARSELT_SPARSITY_50_PERCENT))
+}
+
+// Several things are being done in this stage, algorithm selection, planning,
+// and returning workspace and compressed matrices data buffer sizes.
+extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
+mgpuCuSparseLtSpMMBufferSize(void *workspace_size, void *compressed_size,
+                             void *compressed_buffer_size, void *h, void *a,
+                             CUstream /*stream*/) {
+  // TODO: support more advanced settings, e.g., the input right operand is a
+  // sparse matrix assuming matA is the sparse matrix
+  auto handle = reinterpret_cast<cusparseLtHandle_t *>(h);
+  auto matA = reinterpret_cast<cusparseLtSpMatHandleAndData *>(a);
+
+  CHECK_CUSPARSE(cusparseLtMatmulAlgSelectionInit(
+      handle, &(matWithData.alg_sel), &matmul, CUSPARSELT_MATMUL_ALG_DEFAULT))
+  int alg = 0;
+  CHECK_CUSPARSE(cusparseLtMatmulAlgSetAttribute(
+      handle, &(matWithData.alg_sel), CUSPARSELT_MATMUL_ALG_CONFIG_ID, &alg,
+      sizeof(alg)))
+  // TODO: add transpose support
+  CHECK_CUSPARSE(cusparseLtMatmulDescriptorInit(
+      handle, &(matA.matmul), c, CUSPARSE_OPERATION_NON_TRANSPOSE, &(matA->mat),
+      &matB, &matC, &matC, compute_type))
+  CHECK_CUSPARSE(cusparseLtMatmulPlanInit(handle, &(matWithData.plan), &matmul,
+                                          &(matWithData.alg_sel)))
+
+  CHECK_CUSPARSE(
+      cusparseLtMatmulGetWorkspace(handle, &(matA.plan), workspace_size))
+  CHECK_CUSPARSE(cusparseLtSpMMACompressedSize(
+      handle, &(matA.plan), compressed_size, compressed_buffer_size))
+
+  // avoid zero-alloc
+  *workspace_size = (*workspace_size == 0 ? 1 : *workspace_size);
+  *compressed_size = (*compressed_size == 0 ? 1 : *compressed_size);
+  *compressed_buffer_size =
+      (*compressed_buffer_size == 0 ? 1 : *compressed_buffer_size);
+  return;
+}
+
+extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
+mgpuCuSparseLtSpMM(void *alg_sel, void *plan, void *matmul, void *h, void *a,
+                   void *b, void *c, int32_t dw, void *buf, void *dA_compressed,
+                   void *dA_compressedBuffer, CUstream stream) {
+  auto handle = reinterpret_cast<cusparseLtHandle_t *>(h);
+  auto matA = reinterpret_cast<cusparseLtSpMatHandleAndData *>(a);
+  auto matB = reinterpret_cast<cusparseLtDnMatHandleAndData *>(b);
+  auto matC = reinterpret_cast<cusparseLtDnMatHandleAndData *>(c);
+
+  cusparseLtMatmulAlgSelection_t alg_sel;
+  cusparseLtMatmulPlan_t plan;
+  cusparseLtMatmulDescriptor_t matmul;
+
+  ALPHABETA(dw, alpha, beta)
+
+  CHECK_CUSPARSE(cusparseLtSpMMACompress(handle, &(matA->plan), &(matA->values),
+                                         dA_compressed, dA_compressedBuffer,
+                                         stream))
+
+  // TODO: add support to multi-stream execution
+  // Perform the matrix multiplication. D = A*B+C using C==D for now
+  CHECK_CUSPARSE(
+      cusparseLtMatmul(handle, reinterpret_cast<cusparseLtMatmulPlan_t *>(plan),
+                       &alpha, dA_compressed, dB, &beta, matC->values,
+                       /*dD*/ matC->values, d_workspace, &stream, 1))
+
+  CUSPARSE_REPORT_IF_ERROR(cusparseLtMatDescriptorDestroy(&(mat->mat)))
+  // destroy the plan associated with the sparse matrix
+  CUSPARSE_REPORT_IF_ERROR(cusparseLtMatmulPlanDestroy(&(mat->plan)))
+}
+
+#endif // MLIR_ENABLE_CUDA_CUSPARSELT

mlir/test/Conversion/GPUCommon/2To4Sparsity/lit.local.cfg

+if not config.enable_cuda_runner or not config.mlir_run_cuda_sm80_tests:
+    config.unsupported = True

mlir/test/Conversion/GPUCommon/2To4Sparsity/lower-2to4-sparse-to-gpu-runtime-calls.mlir

+// RUN: mlir-opt %s --gpu-to-llvm='use-opaque-pointers=1' | FileCheck %s
+
+module attributes {gpu.container_module} {
+
+  // CHECK-LABEL: func @matmul
+  // CHECK: llvm.call @mgpuStreamCreate
+  // CHECK: llvm.call @mgpuMemAlloc
+  // CHECK: llvm.call @mgpuMemAlloc
+  // CHECK: llvm.call @mgpuCreateSparseLtEnv
+  // CHECK: llvm.call @mgpuDestroyCuSparseLtSpMat
+  // CHECK: llvm.call @mgpuCreateCuSparseLtDnMat
+  // CHECK: llvm.call @mgpuCuSparseLtSpMMBufferSize
+  // CHECK: llvm.call @mgpuCuSparseLtSpMM
+  // CHECK: llvm.call @mgpuDestroyCuSparseLtSpMat
+  // CHECK: llvm.call @mgpuDestroyCuSparseLtDnMat
+  // CHECK: llvm.call @mgpuDestroySparseLtEnv
+  // CHECK: llvm.call @mgpuStreamSynchronize
+  // CHECK: llvm.call @mgpuStreamDestroy
+  func.func @matmul(%arg0: index) {
+    %token0 = gpu.wait async
+    %mem1, %token1 = gpu.alloc async [%token0] (%arg0) : memref<?xf64>
+    %mem2, %token2 = gpu.alloc async [%token1] (%arg0) : memref<?xf64>
+    %env, %token3 = gpu.create_sparse_env async [%token2]
+    %spmat, %token4 = gpu.create_2to4_spmat async [%token3] %env, %arg0, %arg0, %mem1:  memref<?xf64>
+    %dnmat, %token5 = gpu.create_dn_mat async [%token4] %env, %arg0, %arg0, %mem2 : memref<?xf64>
+    %bufferSzs, %token6 = gpu.spmm_buffer_size async [%token5] %env, %spmat, %dnmat, %dnmat : tuple<index,index,index>
+    %token7 = gpu.spmm async [%token6] %env, %spmat, %dnmat, %dnmat, %mem2, %mem2, %mem2 : memref<?xf64>,memref<?xf64>,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
+  }
+
+}

mlir/test/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir

module attributes {gpu.container_module} {
   // CHECK: llvm.call @mgpuStreamSynchronize
   // CHECK: llvm.call @mgpuStreamDestroy
   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>
+    %dnvec, %token5 = gpu.create_dn_vec async [%token4] %env, %mem2, %arg0 : memref<?xf64>
     %bufferSz, %token6 = gpu.spmv_buffer_size async [%token5] %env, %spmat, %dnvec, %dnvec  into f64
     %token7 = gpu.spmv async [%token6] %env, %spmat, %dnvec, %dnvec, %mem2 : memref<?xf64> into f64
     %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
   }
module attributes {gpu.container_module} {
   // CHECK: llvm.call @mgpuStreamSynchronize
   // CHECK: llvm.call @mgpuStreamDestroy
   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 into f64
+    %dnmat, %token5 = gpu.create_dn_mat async [%token4] %env, %arg0, %arg0, %mem2 : memref<?xf64>
+    %bufferSz, %token6 = gpu.spmm_buffer_size async [%token5] %env, %spmat, %dnmat, %dnmat : index into f64
     %token7 = gpu.spmm async [%token6] %env, %spmat, %dnmat, %dnmat, %mem2 : memref<?xf64> into 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
   }
 
module attributes {gpu.container_module} {
   // CHECK: llvm.call @mgpuStreamSynchronize
   // CHECK: llvm.call @mgpuStreamDestroy
   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>
+    %dnmat, %token5 = gpu.create_dn_mat async [%token4] %env, %arg0, %arg0, %mem2 : memref<?xf64>
     %bufferSz, %token6 = gpu.sddmm_buffer_size async [%token5] %env, %dnmat, %dnmat, %spmat into f64
     %token7 = gpu.sddmm async [%token6] %env, %dnmat, %dnmat, %spmat, %mem2 : memref<?xf64> into 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
   }
 
 }
 
 

mlir/test/Dialect/GPU/ops.mlir

func.func @sparse_ops(%arg0: index) {
     %mem2, %token2 = gpu.alloc async [%token1] (%arg0) : memref<?xf64>
     // CHECK: gpu.create_sparse_env async
     %env, %token3 = gpu.create_sparse_env async [%token2]
     // CHECK: gpu.create_coo async
     %spmat, %token4 = gpu.create_coo async [%token3] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
     // CHECK: gpu.create_csr async
     %spmat2, %token5 = gpu.create_csr async [%token4] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
     // CHECK: gpu.create_dn_vec async
-    %dnvec, %token6 = gpu.create_dn_vec async [%token5] %mem2, %arg0 : memref<?xf64>
+    %dnvec, %token6 = gpu.create_dn_vec async [%token5] %env, %mem2, %arg0 : memref<?xf64>
     // CHECK: gpu.spmv_buffer_size async
     %bufferSz, %token7 = gpu.spmv_buffer_size async [%token6] %env, %spmat, %dnvec, %dnvec  into f64
     // CHECK: gpu.spmv async
     %token8 = gpu.spmv async [%token7] %env, %spmat, %dnvec, %dnvec, %mem2 : memref<?xf64>  into f64
     // CHECK: gpu.create_dn_mat async
-    %dnmat, %token9 = gpu.create_dn_mat async [%token8] %arg0, %arg0, %mem2 : memref<?xf64>
+    %dnmat, %token9 = gpu.create_dn_mat async [%token8] %env, %arg0, %arg0, %mem2 : memref<?xf64>
     // CHECK: gpu.spmm_buffer_size async
-    %bufferSz2, %token10 = gpu.spmm_buffer_size async [%token9] %env, %spmat, %dnmat, %dnmat  into f64
+    %bufferSz2, %token10 = gpu.spmm_buffer_size async [%token9] %env, %spmat, %dnmat, %dnmat : index into f64
     // CHECK: gpu.spmm async
     %token11 = gpu.spmm async [%token10] %env, %spmat, %dnmat, %dnmat, %mem2 : memref<?xf64>  into f64
     // CHECK: gpu.sddmm_buffer_size async
     %bufferSz3, %token12 = gpu.sddmm_buffer_size async [%token11] %env, %dnmat, %dnmat, %spmat  into f64
     // CHECK: gpu.sddmm async
     %token13 = gpu.sddmm async [%token12] %env, %dnmat, %dnmat, %spmat, %mem2 : memref<?xf64>  into f64
     // CHECK: gpu.destroy_dn_mat async
     %token14 = gpu.destroy_dn_mat async [%token13] %dnmat

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.create_dn_vec async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}} : memref<?xf64>
   // CHECK: %{{.*}}, %{{.*}} = gpu.spmv_buffer_size async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} into f64
   // CHECK: %{{.*}} = gpu.spmv async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xf64> into f64
   // 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>
+    %dnvec, %token5 = gpu.create_dn_vec async [%token4] %env, %mem2, %arg0 : memref<?xf64>
     %bufferSz, %token6 = gpu.spmv_buffer_size async [%token5] %env, %spmat, %dnvec, %dnvec into f64
     %token7 = gpu.spmv async [%token6] %env, %spmat, %dnvec, %dnvec, %mem2 : memref<?xf64> into f64
     %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.create_dn_mat async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xf64>
   // CHECK: %{{.*}}, %{{.*}} = gpu.spmm_buffer_size async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} into f64
   // CHECK: %{{.*}} = gpu.spmm async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xf64> into 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 into f64
+    %dnmat, %token5 = gpu.create_dn_mat async [%token4] %env, %arg0, %arg0, %mem2 : memref<?xf64>
+    %bufferSz, %token6 = gpu.spmm_buffer_size async [%token5] %env, %spmat, %dnmat, %dnmat : index into f64
     %token7 = gpu.spmm async [%token6] %env, %spmat, %dnmat, %dnmat, %mem2 : memref<?xf64> into 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.create_dn_mat async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xf64>
   // CHECK: %{{.*}}, %{{.*}} = gpu.sddmm_buffer_size async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}  into f64
   // CHECK: %{{.*}} = gpu.sddmm async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xf64>  into 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 @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>
+    %dnmat, %token5 = gpu.create_dn_mat async [%token4] %env, %arg0, %arg0, %mem2 : memref<?xf64>
     %bufferSz, %token6 = gpu.sddmm_buffer_size async [%token5] %env, %dnmat, %dnmat, %spmat into f64
     %token7 = gpu.sddmm async [%token6] %env, %dnmat, %dnmat, %spmat, %mem2 : memref<?xf64> into 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
   }
 
 }
 
 

mlir/test/Dialect/SparseTensor/GPU/gpu_matmul_lib.mlir

 // CHECK:           %[[VAL_36:.*]] = memref.dim %[[VAL_34]], %[[VAL_3]] : memref<?x?xf64>
 // CHECK:           %[[VAL_37:.*]] = memref.dim %[[VAL_34]], %[[VAL_4]] : memref<?x?xf64>
 // CHECK:           %[[VAL_38:.*]], %[[VAL_39:.*]] = gpu.alloc async {{\[}}%[[VAL_35]]] (%[[VAL_36]], %[[VAL_37]]) : memref<?x?xf64>
 // CHECK:           %[[VAL_40:.*]] = gpu.memcpy async {{\[}}%[[VAL_39]]] %[[VAL_38]], %[[VAL_34]] : memref<?x?xf64>, memref<?x?xf64>
 // CHECK:           gpu.wait {{\[}}%[[VAL_16]], %[[VAL_21]], %[[VAL_26]], %[[VAL_33]], %[[VAL_40]]]
 // CHECK:           %[[VAL_41:.*]] = gpu.wait async
 // CHECK:           %[[VAL_42:.*]], %[[VAL_43:.*]] = gpu.create_sparse_env async {{\[}}%[[VAL_41]]]
 // CHECK:           %[[VAL_44:.*]], %[[VAL_45:.*]] = gpu.create_csr async {{\[}}%[[VAL_43]]] %[[VAL_6]], %[[VAL_7]], %[[VAL_5]], %[[VAL_14]], %[[VAL_19]], %[[VAL_24]] : memref<?xindex>, memref<?xindex>, memref<?xf64>
-// CHECK:           %[[VAL_46:.*]], %[[VAL_47:.*]] = gpu.create_dn_mat async {{\[}}%[[VAL_45]]] %[[VAL_7]], %[[VAL_8]], %[[VAL_31]] : memref<?x?xf64>
-// CHECK:           %[[VAL_48:.*]], %[[VAL_49:.*]] = gpu.create_dn_mat async {{\[}}%[[VAL_47]]] %[[VAL_6]], %[[VAL_8]], %[[VAL_38]] : memref<?x?xf64>
-// CHECK:           %[[VAL_50:.*]], %[[VAL_51:.*]] = gpu.spmm_buffer_size async {{\[}}%[[VAL_49]]] %[[VAL_42]], %[[VAL_44]], %[[VAL_46]], %[[VAL_48]]
+// CHECK:           %[[VAL_46:.*]], %[[VAL_47:.*]] = gpu.create_dn_mat async {{\[}}%[[VAL_45]]] %[[VAL_42]], %[[VAL_7]], %[[VAL_8]], %[[VAL_31]] : memref<?x?xf64>
+// CHECK:           %[[VAL_48:.*]], %[[VAL_49:.*]] = gpu.create_dn_mat async {{\[}}%[[VAL_47]]] %[[VAL_42]], %[[VAL_6]], %[[VAL_8]], %[[VAL_38]] : memref<?x?xf64>
+// CHECK:           %[[VAL_50:.*]], %[[VAL_51:.*]] = gpu.spmm_buffer_size async {{\[}}%[[VAL_49]]] %[[VAL_42]], %[[VAL_44]], %[[VAL_46]], %[[VAL_48]] : index
 // CHECK:           %[[VAL_52:.*]], %[[VAL_53:.*]] = gpu.alloc async {{\[}}%[[VAL_51]]] (%[[VAL_50]]) : memref<?xi8>
 // CHECK:           %[[VAL_54:.*]] = gpu.spmm async {{\[}}%[[VAL_53]]] %[[VAL_42]], %[[VAL_44]], %[[VAL_46]], %[[VAL_48]], %[[VAL_52]] : memref<?xi8>
 // CHECK:           %[[VAL_55:.*]] = gpu.destroy_sp_mat async {{\[}}%[[VAL_54]]] %[[VAL_44]]
 // CHECK:           %[[VAL_56:.*]] = gpu.destroy_dn_mat async {{\[}}%[[VAL_55]]] %[[VAL_46]]
 // CHECK:           %[[VAL_57:.*]] = gpu.destroy_dn_mat async {{\[}}%[[VAL_56]]] %[[VAL_48]]
 // CHECK:           %[[VAL_58:.*]] = gpu.destroy_sparse_env async {{\[}}%[[VAL_57]]] %[[VAL_42]]
 // CHECK:           %[[VAL_59:.*]] = gpu.dealloc async {{\[}}%[[VAL_58]]] %[[VAL_14]] : memref<?xindex>
 // CHECK:           %[[VAL_60:.*]] = gpu.dealloc async {{\[}}%[[VAL_59]]] %[[VAL_19]] : memref<?xindex>

mlir/test/Dialect/SparseTensor/GPU/gpu_matvec_lib.mlir

 // CHECK:           %[[VAL_33:.*]] = gpu.wait async
 // CHECK:           %[[VAL_34:.*]] = memref.dim %[[VAL_32]], %[[VAL_3]] : memref<?xf64>
 // CHECK:           %[[VAL_35:.*]], %[[VAL_36:.*]] = gpu.alloc async {{\[}}%[[VAL_33]]] (%[[VAL_34]]) : memref<?xf64>
 // CHECK:           %[[VAL_37:.*]] = gpu.memcpy async {{\[}}%[[VAL_36]]] %[[VAL_35]], %[[VAL_32]] : memref<?xf64>, memref<?xf64>
 // CHECK:           gpu.wait {{\[}}%[[VAL_15]], %[[VAL_20]], %[[VAL_25]], %[[VAL_31]], %[[VAL_37]]]
 // CHECK:           %[[VAL_38:.*]] = gpu.wait async
 // CHECK:           %[[VAL_39:.*]], %[[VAL_40:.*]] = gpu.create_sparse_env async {{\[}}%[[VAL_38]]]
 // CHECK:           %[[VAL_41:.*]], %[[VAL_42:.*]] = gpu.create_coo async {{\[}}%[[VAL_40]]] %[[VAL_6]], %[[VAL_7]], %[[VAL_5]], %[[VAL_13]], %[[VAL_18]], %[[VAL_23]] : memref<?xindex>, memref<?xindex>, memref<?xf64>
-// CHECK:           %[[VAL_43:.*]], %[[VAL_44:.*]] = gpu.create_dn_vec async {{\[}}%[[VAL_42]]] %[[VAL_29]], %[[VAL_7]] : memref<?xf64>
-// CHECK:           %[[VAL_45:.*]], %[[VAL_46:.*]] = gpu.create_dn_vec async {{\[}}%[[VAL_44]]] %[[VAL_35]], %[[VAL_6]] : memref<?xf64>
+// CHECK:           %[[VAL_43:.*]], %[[VAL_44:.*]] = gpu.create_dn_vec async {{\[}}%[[VAL_42]]] %[[VAL_39:.*]], %[[VAL_29]], %[[VAL_7]] : memref<?xf64>
+// CHECK:           %[[VAL_45:.*]], %[[VAL_46:.*]] = gpu.create_dn_vec async {{\[}}%[[VAL_44]]] %[[VAL_39:.*]], %[[VAL_35]], %[[VAL_6]] : memref<?xf64>
 // CHECK:           %[[VAL_47:.*]], %[[VAL_48:.*]] = gpu.spmv_buffer_size async {{\[}}%[[VAL_46]]] %[[VAL_39]], %[[VAL_41]], %[[VAL_43]], %[[VAL_45]]
 // CHECK:           %[[VAL_49:.*]], %[[VAL_50:.*]] = gpu.alloc async {{\[}}%[[VAL_48]]] (%[[VAL_47]]) : memref<?xi8>
 // CHECK:           %[[VAL_51:.*]] = gpu.spmv async {{\[}}%[[VAL_50]]] %[[VAL_39]], %[[VAL_41]], %[[VAL_43]], %[[VAL_45]], %[[VAL_49]] : memref<?xi8>
 // CHECK:           %[[VAL_52:.*]] = gpu.destroy_sp_mat async {{\[}}%[[VAL_51]]] %[[VAL_41]]
 // CHECK:           %[[VAL_53:.*]] = gpu.destroy_dn_vec async {{\[}}%[[VAL_52]]] %[[VAL_43]]
 // CHECK:           %[[VAL_54:.*]] = gpu.destroy_dn_vec async {{\[}}%[[VAL_53]]] %[[VAL_45]]
 // CHECK:           %[[VAL_55:.*]] = gpu.destroy_sparse_env async {{\[}}%[[VAL_54]]] %[[VAL_39]]
 // CHECK:           %[[VAL_56:.*]] = gpu.dealloc async {{\[}}%[[VAL_55]]] %[[VAL_13]] : memref<?xindex>