[mlir] [gpu] [sparse] refined SparseHandle type

[aartbik]

please rebase against main
(since these files were extended)

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

114

do we want to refine the "sparse handle type" strings as well for all four now?

mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h

171

move this up to other includes

also, always double check that you don't need changes in bazel/cmake when you add a dependence

174

for my understanding, why default? can we make e.g. Env default ;-)

181

use same order as above, i.e. env first

[K-Wu]

upd gpubase.td

[K-Wu]

upd class names in the commits rebase-pulled from origin/main

[K-Wu]

fix broken tests

[K-Wu]

I have fixed this diff so that all tests are passed

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

114

good catch! I added a few words here

mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h

171

Good catch! I tested and found it was not needed. I added that during the debugging process.

174

That makes sense. I remove that

181

Done

[aartbik]

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

554

update here and elsewhere in SpMM

[K-Wu]

update to refined sparse handles in incoming changes

[K-Wu]

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

554

Got it! This is updated now

[aartbik]

mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h

189

more modern casting?

[aartbik]

mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h

188

intentional?

[K-Wu]

rm the unnecessary classof impl

[K-Wu]

mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h

189

This custom classof is completely unnecessary. credit to Peiming and wrengr

[K-Wu]

rm unnecessary code

[K-Wu]

readability

[wrengr]

mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h

173–174

You need to add the "typename" keyword after the equals sign. This will fix some of the build errors.

(For the other build errors, be sure to rerun git clang-format HEAD^ and then upload the new version of the differential)

[wrengr]

Just a few style remarks, but other than that this is looking pretty good

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

116

Could you reformat this section to:

• Add a blank line between each definition
• Reflow the code to avoid going over 80 columns. E.g., using the following format:

def GPU_SparseDnVecHandle :
  DialectType<GPU_Dialect,
    CPred<"$_self.isa<::mlir::gpu::SparseDnVecHandleType>()">,
    "sparse dense vector handle type">,
  BuildableType<"mlir::gpu::SparseDnVecHandleType::get($_builder.getContext())">;

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

1560–1561

Nit: Here and below, can you reformat the results line to avoid going over 80 columns? Indentation should match the other let arguments = (ins ...); and let results = (outs ...); in this file

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1484–1485

Why not define the helper template-function I suggested over chat, so as to avoid repeating so much code?

(In case you forgot, that was:

template<typename T>
static void addOpaquePointerConversion(LLVMTypeConverter &converter) {
  converter.addConversion([&converter](T) -> Type {
    return converter.getPointerType(
      IntegerType::get(&converter.getContext(), 8));
  });
}

Then the definition of this function becomes just:

void mlir::populateGpuToLLVMConversionPatterns(LLVMTypeConverter &converter, .....) {
  addOpaquePointerConversion<gpu::AsyncTokenType>(converter);
  addOpaquePointerConversion<gpu::SparseDnVecHandleType>(converter);
  addOpaquePointerConversion<gpu::SparseDnMatHandleType>(converter);
  .....

[K-Wu]

In D151014#4363054, @wrengr wrote:

Just a few style remarks, but other than that this is looking pretty good

I will address these two comments tomorrow.

[K-Wu]

reformat and changes to make sure build could pass

[K-Wu]

add helper func as wrengr suggested

[K-Wu]

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

116

Done

mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h

173–174

Now I know I need to check the build log. Thanks!

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

1560–1561

Worked on it

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1484–1485

I forgot that. THank you for this!

[K-Wu]

reformat

[K-Wu]

fixe variable name issue when rebase-pulling

[aartbik]

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

122

nit: "sparse dense vector handle" reads a bit weird (I see why you picked this, but then :134 should really be "sparse sparse matrix handle" for consistency). I would just remove the "sparse" for DnVec and DnMat from the description string

134

remove tailing "in sparse"

mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h

178

this is nice!

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

150

same order as in decl, so spmat last

222

same order as decl please, so env/dnvec/dnmat/spmat

232–233

it would probably be good to have a central place for this string (in the type), so that printing and parsing matches

238

here too

[K-Wu]

addressing the order of refined HandleType

[K-Wu]

mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h

178

credit to

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

232–233

How about something like adding a static constexpr string in each HandleType definition, i.e., those template specialization lines?

[K-Wu]

mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h

178

Great thanks to @wrengr and @Peiming

[K-Wu]

centralize keywords

[K-Wu]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

232–233

It turned out this is way too complicated and static variables needs to be defined in .cpp anyway.
I am using std::map to centralize the keywords

[aartbik]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

232–233

but can it be static? we are not in an anonymous namespace at this point?
or move it there?

[Peiming]

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

1705

Here and belowShould be aligned as others change you made.

[K-Wu]

add static keyword

[K-Wu]

formatted the td

[K-Wu]

fix typo

[K-Wu]

lazy eval to fix global destructor compiler warning

[aartbik]

rebased with main

[NathanielMcVicar]

This appears to have broken the Windows bot: https://lab.llvm.org/buildbot/#/builders/13/builds/35754/steps/6/logs/stdio

C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\lib\Dialect\GPU\IR\GPUDialect.cpp(175) : error C2220: the following warning is treated as an error
C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\lib\Dialect\GPU\IR\GPUDialect.cpp(175) : warning C4715: 'getSparseHandleKeyword': not all control paths return a value

[K-Wu]

@NathanielMcVicar Thank you for catching this! I am creating a new diff to fix this issue https://reviews.llvm.org/D151405

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

 // Memref type acceptable to gpu.subgroup_mma_{load|store}_matrix ops.
 def GPU_MMAMemRef : MemRefOf<[I8, I32, F16, F32, VectorOfRankAndType<[1], [I8, I32, F16, F32]>]>;
 
 class MMAMatrixOf<list<Type> allowedTypes> :
   ContainerType<AnyTypeOf<allowedTypes>, IsMMAMatrixTypePred,
   "::llvm::cast<::mlir::gpu::MMAMatrixType>($_self).getElementType()",
   "gpu.mma_matrix", "::mlir::gpu::MMAMatrixType">;
 
-// Generic type for all sparse handles (could be refined).
-def GPU_SparseHandle : DialectType<
-  GPU_Dialect, CPred<"$_self.isa<::mlir::gpu::SparseHandleType>()">, "sparse handle type">,
-             BuildableType<"mlir::gpu::SparseHandleType::get($_builder.getContext())">;
+// Types for all sparse handles.
+def GPU_SparseEnvHandle : 
+  DialectType<GPU_Dialect, 
+    CPred<"$_self.isa<::mlir::gpu::SparseEnvHandleType>()">, 
+    "sparse environment handle type">, 
+  BuildableType<"mlir::gpu::SparseEnvHandleType::get($_builder.getContext())">;
+
+def GPU_SparseDnVecHandle : 
+  DialectType<GPU_Dialect, 
+    CPred<"$_self.isa<::mlir::gpu::SparseDnVecHandleType>()">, 
+    "dense vector handle type">,
+  BuildableType<"mlir::gpu::SparseDnVecHandleType::get($_builder.getContext())">;
+
+def GPU_SparseDnMatHandle : 
+  DialectType<GPU_Dialect, 
+    CPred<"$_self.isa<::mlir::gpu::SparseDnMatHandleType>()">, 
+    "dense matrix handle type">,
+  BuildableType<"mlir::gpu::SparseDnMatHandleType::get($_builder.getContext())">;
+
+def GPU_SparseSpMatHandle : 
+  DialectType<GPU_Dialect, 
+    CPred<"$_self.isa<::mlir::gpu::SparseSpMatHandleType>()">, 
+    "sparse matrix handle type">,
+  BuildableType<"mlir::gpu::SparseSpMatHandleType::get($_builder.getContext())">;
 
 //===----------------------------------------------------------------------===//
 // GPU Interfaces.
 //===----------------------------------------------------------------------===//
 
 def GPU_AsyncOpInterface : OpInterface<"AsyncOpInterface"> {
   let description = [{
     Interface for GPU operations that execute asynchronously on the device.

mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h

public:
   /// C += A*B. This function returns which operand in the given equation is
   /// held by this type. String returned can be one of"AOp", "BOp" and "COp".
   StringRef getOperand() const;
 };
 
 // Adds a `gpu.async.token` to the front of the argument list.
 void addAsyncDependency(Operation *op, Value token);
 
-// Represents any sparse handle.
+// Handle types for sparse.
+enum class SparseHandleKind { Env, DnVec, DnMat, SpMat };
+
+template <SparseHandleKind K>
 class SparseHandleType
-    : public Type::TypeBase<SparseHandleType, Type, TypeStorage> {
+    : public Type::TypeBase<SparseHandleType<K>, Type, TypeStorage> {
 public:
-  // Used for generic hooks in TypeBase.
+  using Base =
+      typename Type::TypeBase<SparseHandleType<K>, Type, TypeStorage>::Base;
   using Base::Base;
 };
 
+using SparseEnvHandleType = SparseHandleType<SparseHandleKind::Env>;
+using SparseDnVecHandleType = SparseHandleType<SparseHandleKind::DnVec>;
+using SparseDnMatHandleType = SparseHandleType<SparseHandleKind::DnMat>;
+using SparseSpMatHandleType = SparseHandleType<SparseHandleKind::SpMat>;
+
 } // namespace gpu
 } // namespace mlir
 
 #include "mlir/Dialect/GPU/IR/GPUOpsEnums.h.inc"
 
 #include "mlir/Dialect/GPU/IR/GPUOpsDialect.h.inc"
 
 #include "mlir/Dialect/GPU/IR/GPUOpInterfaces.h.inc"
 
 #include "mlir/Dialect/SCF/IR/DeviceMappingInterface.h"
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/GPU/IR/GPUOpsAttributes.h.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/GPU/IR/GPUOps.h.inc"
 
 #endif // MLIR_DIALECT_GPU_IR_GPUDIALECT_H

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

let description = [{
     Example:
 
     ```mlir
     %env, %token = gpu.create_sparse_env async [%dep]
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies);
-  let results = (outs Res<GPU_SparseHandle>:$env, Optional<GPU_AsyncToken>:$asyncToken);
-
+  let results = (outs Res<GPU_SparseEnvHandle>:$env, 
+                      Optional<GPU_AsyncToken>:$asyncToken);
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) attr-dict
   }];
 }
 
-def GPU_DestroySparseEnvOp : GPU_Op<"destroy_sparse_env", [GPU_AsyncOpInterface]> {
+def GPU_DestroySparseEnvOp : GPU_Op<
+    "destroy_sparse_env", 
+    [GPU_AsyncOpInterface]> {
   let summary = "Destroy sparse environment operation";
   let description = [{
     The `gpu.destroy_sparse_env` operation releases all resources of a sparse
     environment represented by a handle that was previously created by a
     `gpu.create_sparse_env` operation.
 
     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
     %token = gpu.destroy_sparse_env async [%dep] %env
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
-                   Arg<GPU_SparseHandle>:$env);
+                       Arg<GPU_SparseEnvHandle>:$env);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
-    custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) $env attr-dict
+    custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) 
+    $env attr-dict
   }];
 }
 
 def GPU_CreateDnVecOp : GPU_Op<"create_dn_vec", [GPU_AsyncOpInterface]> {
   let summary = "Create dense vector operation";
   let description = [{
     The `gpu.create_dn_vec` operation initializes a dense vector from
     the given values buffer and size. The buffer must already be copied
     from the host to the device prior to using this operation. The
     operation returns a handle to the dense vector 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
     %dvec, %token = gpu.create_dn_vec async [%dep] %mem, %size : memref<?xf64>
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                    AnyMemRef:$memref, Index:$size);
-  let results = (outs Res<GPU_SparseHandle>:$dvec, Optional<GPU_AsyncToken>:$asyncToken);
+  let results = (outs Res<GPU_SparseDnVecHandle>:$dvec, 
+                      Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $memref `,` $size attr-dict `:` type($memref)
   }];
 }
 
 def GPU_DestroyDnVecOp : GPU_Op<"destroy_dn_vec", [GPU_AsyncOpInterface]> {
let description = [{
     Example:
 
     ```mlir
     %token = gpu.destroy_dn_vec async [%dep] %dvec
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
-                   Arg<GPU_SparseHandle>:$dvec);
+                       Arg<GPU_SparseDnVecHandle>:$dvec);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
-    custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) $dvec attr-dict
+    custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) 
+    $dvec attr-dict
   }];
 }
 
 def GPU_CreateDnMatOp : GPU_Op<"create_dn_mat", [GPU_AsyncOpInterface]> {
   let summary = "Create dense matrix operation";
   let description = [{
     The `gpu.create_dn_mat` operation initializes a dense matrix from
     the given values buffer and sizes. The buffer must already be copied
     from the host to the device prior to using this operation. The
     operation returns a handle to the dense 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
     %dmat, %token = gpu.create_dn_mat async [%dep] %mem, %size : memref<?xf64>
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                        Index:$rows,
                        Index:$cols,
                        AnyMemRef:$memref);
-  let results = (outs Res<GPU_SparseHandle>:$dmat, Optional<GPU_AsyncToken>:$asyncToken);
+  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)
   }];
 }
 
 def GPU_DestroyDnMatOp : GPU_Op<"destroy_dn_mat", [GPU_AsyncOpInterface]> {
let description = [{
     Example:
 
     ```mlir
     %token = gpu.destroy_dn_vec async [%dep] %dmat
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
-                   Arg<GPU_SparseHandle>:$dmat);
+                       Arg<GPU_SparseDnMatHandle>:$dmat);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
-    custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) $dmat attr-dict
+    custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) 
+    $dmat attr-dict
   }];
 }
 
 def GPU_CreateCooOp : GPU_Op<"create_coo", [GPU_AsyncOpInterface]> {
   let summary = "Create sparse matrix in COO format operation";
   let description = [{
     The `gpu.create_coo` operation initializes a sparse matrix in COO format
     with the given sizes from the given index and values buffers. 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_coo async [%dep] %rows, %cols, %nnz, %rowIdx,
         %colIdx, %values : memref<?xindex>, memref<?xindex>, memref<?xf64>
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                        Index:$rows,
                        Index:$cols,
                        Index:$nnz,
                        AnyMemRef:$rowIdxs,
                        AnyMemRef:$colIdxs,
                        AnyMemRef:$values);
-  let results = (outs Res<GPU_SparseHandle>:$spmat, Optional<GPU_AsyncToken>:$asyncToken);
+  let results = (outs Res<GPU_SparseSpMatHandle>:$spmat, 
+                      Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $rows `,` $cols `,` $nnz `,` $rowIdxs `,` $colIdxs `,` $values attr-dict
     `:` type($rowIdxs) `,` type($colIdxs) `,` type($values)
   }];
 }
 
def GPU_CreateCsrOp : GPU_Op<"create_csr", [GPU_AsyncOpInterface]> {
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                    Index:$rows,
                    Index:$cols,
                    Index:$nnz,
                    AnyMemRef:$rowPos,
                    AnyMemRef:$colIdxs,
                    AnyMemRef:$values);
-  let results = (outs Res<GPU_SparseHandle>:$spmat, Optional<GPU_AsyncToken>:$asyncToken);
+  let results = (outs Res<GPU_SparseSpMatHandle>:$spmat, 
+                      Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $rows `,` $cols `,` $nnz `,` $rowPos `,` $colIdxs `,` $values attr-dict
     `:` type($rowPos) `,` type($colIdxs) `,` type($values)
   }];
 }
 
let description = [{
     Example:
 
     ```mlir
     %token = gpu.destroy_sp_mat async [%dep] %spmat
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
-                   Arg<GPU_SparseHandle>:$spmat);
+                       Arg<GPU_SparseSpMatHandle>:$spmat);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) $spmat attr-dict
   }];
 }
 
 def GPU_SpMVBufferSizeOp : GPU_Op<"spmv_buffer_size", [GPU_AsyncOpInterface]> {
let description = [{
     that case, it returns a !gpu.async.token in addition to the environment.
 
     Example:
 
     ```mlir
     %buffersz, %token = gpu.spmv_buffersize async [%dep] %env, %spmatA, %dnX, %dnY
     ```
   }];
-
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
-                   GPU_SparseHandle:$env,
-                   GPU_SparseHandle:$spmatA,
-                   GPU_SparseHandle:$dnX,
-                   GPU_SparseHandle:$dnY);
-  let results = (outs Res<Index>:$bufferSz, Optional<GPU_AsyncToken>:$asyncToken);
+                       GPU_SparseEnvHandle:$env,
+                       GPU_SparseSpMatHandle:$spmatA,
+                       GPU_SparseDnVecHandle:$dnX,
+                       GPU_SparseDnVecHandle:$dnY);
+  let results = (outs Res<Index>:$bufferSz, 
+                      Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $env `,` $spmatA `,` $dnX `,` $dnY attr-dict
   }];
 }
 
 def GPU_SpMVOp : GPU_Op<"spmv", [GPU_AsyncOpInterface]> {
let description = [{
     that case, it returns a !gpu.async.token in addition to the environment.
 
     Example:
 
     ```mlir
     %token = gpu.spmv async [%dep] %env, %spmatA, %dnX, %dnY : memref<?xf64>
     ```
   }];
-
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
-                   GPU_SparseHandle:$env,
-                   GPU_SparseHandle:$spmatA,
-                   GPU_SparseHandle:$dnX,
-                   GPU_SparseHandle:$dnY,
+                       GPU_SparseEnvHandle:$env,
+                       GPU_SparseSpMatHandle:$spmatA,
+                       GPU_SparseDnVecHandle:$dnX,
+                       GPU_SparseDnVecHandle:$dnY,
                        AnyMemRef:$buffer);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $env `,` $spmatA `,` $dnX `,` $dnY `,` $buffer attr-dict `:` type($buffer)
   }];
 }
 
let description = [{
     Example:
 
     ```mlir
     %buffersz, %token = gpu.spmm_buffersize async [%dep] %env, %spmatA, %spmatB, %spmatC
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
-                   GPU_SparseHandle:$env,
-                   GPU_SparseHandle:$spmatA,
-                   GPU_SparseHandle:$dnmatB,
-                   GPU_SparseHandle:$dnmatC);
-  let results = (outs Res<Index>:$bufferSz, Optional<GPU_AsyncToken>:$asyncToken);
+                       GPU_SparseEnvHandle:$env,
+                       GPU_SparseSpMatHandle:$spmatA,
+                       GPU_SparseDnMatHandle:$dnmatB,
+                       GPU_SparseDnMatHandle:$dnmatC);
+  let results = (outs Res<Index>:$bufferSz, 
+                      Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $env `,` $spmatA `,` $dnmatB `,` $dnmatC attr-dict
   }];
 }
 
 def GPU_SpMMOp : GPU_Op<"spmm", [GPU_AsyncOpInterface]> {
let description = [{
     Example:
 
     ```mlir
     %token = gpu.spmm async [%dep] %env, %spmatA, %spmatB, %spmatC, %buffer
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
-                   GPU_SparseHandle:$env,
-                   GPU_SparseHandle:$spmatA,
-                   GPU_SparseHandle:$dnmatB,
-                   GPU_SparseHandle:$dnmatC,
+                       GPU_SparseEnvHandle:$env,
+                       GPU_SparseSpMatHandle:$spmatA,
+                       GPU_SparseDnMatHandle:$dnmatB,
+                       GPU_SparseDnMatHandle:$dnmatC,
                        AnyMemRef:$buffer);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $env `,` $spmatA `,` $dnmatB `,` $dnmatC `,` $buffer attr-dict `:` type($buffer)
   }];
 }
 
 #endif // GPU_OPS

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

LogicalResult ConvertSpMMOpToGpuRuntimeCallPattern::matchAndRewrite(
   spMMCallBuilder.create(loc, rewriter,
                          {adaptor.getEnv(), adaptor.getSpmatA(),
                           adaptor.getDnmatB(), adaptor.getDnmatC(), dw, 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));
+  });
+}
+
 void mlir::populateGpuToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                                RewritePatternSet &patterns,
                                                StringRef gpuBinaryAnnotation,
                                                bool kernelBarePtrCallConv) {
-  converter.addConversion([&converter](gpu::AsyncTokenType type) -> Type {
-    return converter.getPointerType(
-        IntegerType::get(&converter.getContext(), 8));
-  });
-  converter.addConversion([&converter](gpu::SparseHandleType type) -> Type {
-    return converter.getPointerType(
-        IntegerType::get(&converter.getContext(), 8));
-  });
+  addOpaquePointerConversion<gpu::AsyncTokenType>(converter);
+  addOpaquePointerConversion<gpu::SparseDnVecHandleType>(converter);
+  addOpaquePointerConversion<gpu::SparseDnMatHandleType>(converter);
+  addOpaquePointerConversion<gpu::SparseSpMatHandleType>(converter);
+  addOpaquePointerConversion<gpu::SparseEnvHandleType>(converter);
 
   patterns.add<ConvertAllocOpToGpuRuntimeCallPattern,
                ConvertDeallocOpToGpuRuntimeCallPattern,
                ConvertHostRegisterOpToGpuRuntimeCallPattern,
                ConvertHostUnregisterOpToGpuRuntimeCallPattern,
                ConvertMemcpyOpToGpuRuntimeCallPattern,
                ConvertMemsetOpToGpuRuntimeCallPattern,
                ConvertSetDefaultDeviceOpToGpuRuntimeCallPattern,

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
     return true;
   }
 };
 } // namespace
 
 void GPUDialect::initialize() {
   addTypes<AsyncTokenType>();
   addTypes<MMAMatrixType>();
-  addTypes<SparseHandleType>();
+  addTypes<SparseEnvHandleType>();
+  addTypes<SparseDnVecHandleType>();
+  addTypes<SparseDnMatHandleType>();
+  addTypes<SparseSpMatHandleType>();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/GPU/IR/GPUOps.cpp.inc"
       >();
   addAttributes<
 #define GET_ATTRDEF_LIST
 #include "mlir/Dialect/GPU/IR/GPUOpsAttributes.cpp.inc"
       >();
   addInterfaces<GPUInlinerInterface>();
 }
 
+static std::string getSparseHandleKeyword(SparseHandleKind kind) {
+  switch (kind) {
+  case SparseHandleKind::Env:
+    return "sparse.env_handle";
+  case SparseHandleKind::DnVec:
+    return "sparse.dnvec_handle";
+  case SparseHandleKind::DnMat:
+    return "sparse.dnmat_handle";
+  case SparseHandleKind::SpMat:
+    return "sparse.spmat_handle";
+  }
+};
+
 Type GPUDialect::parseType(DialectAsmParser &parser) const {
   // Parse the main keyword for the type.
   StringRef keyword;
   if (parser.parseKeyword(&keyword))
     return Type();
   MLIRContext *context = getContext();
 
   // Handle 'async token' types.
if (keyword == "mma_matrix") {
     if (parser.parseGreater())
       return nullptr;
 
     return MMAMatrixType::getChecked(mlir::detail::getDefaultDiagnosticEmitFn(
                                          parser.getEncodedSourceLoc(beginLoc)),
                                      shape, elementType, operand);
   }
 
-  if (keyword == "sparse.handle")
-    return SparseHandleType::get(context);
+  if (keyword == getSparseHandleKeyword(SparseHandleKind::Env))
+    return SparseEnvHandleType::get(context);
+  if (keyword == getSparseHandleKeyword(SparseHandleKind::DnVec))
+    return SparseDnVecHandleType::get(context);
+  if (keyword == getSparseHandleKeyword(SparseHandleKind::DnMat))
+    return SparseDnMatHandleType::get(context);
+  if (keyword == getSparseHandleKeyword(SparseHandleKind::SpMat))
+    return SparseSpMatHandleType::get(context);
 
   parser.emitError(parser.getNameLoc(), "unknown gpu type: " + keyword);
   return Type();
 }
-
+// TODO: print refined type here. Notice that should be corresponding to the
+// parser
 void GPUDialect::printType(Type type, DialectAsmPrinter &os) const {
   TypeSwitch<Type>(type)
       .Case<AsyncTokenType>([&](Type) { os << "async.token"; })
-      .Case<SparseHandleType>([&](Type) { os << "sparse.handle"; })
+      .Case<SparseEnvHandleType>(
+          [&](Type) { os << getSparseHandleKeyword(SparseHandleKind::Env); })
+      .Case<SparseDnVecHandleType>(
+          [&](Type) { os << getSparseHandleKeyword(SparseHandleKind::DnVec); })
+      .Case<SparseDnMatHandleType>(
+          [&](Type) { os << getSparseHandleKeyword(SparseHandleKind::DnMat); })
+      .Case<SparseSpMatHandleType>(
+          [&](Type) { os << getSparseHandleKeyword(SparseHandleKind::SpMat); })
       .Case<MMAMatrixType>([&](MMAMatrixType fragTy) {
         os << "mma_matrix<";
         auto shape = fragTy.getShape();
         for (auto dim = shape.begin(), e = shape.end() - 1; dim != e; ++dim)
           os << *dim << 'x';
         os << shape.back() << 'x' << fragTy.getElementType();
         os << ", \"" << fragTy.getOperand() << "\"" << '>';
       })

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

static LogicalResult rewriteSpMV(PatternRewriter &rewriter,
   Value vecX = genAllocCopy(rewriter, loc, memX, tokens);
   Value memY = genTensorToMemref(rewriter, loc, y);
   Value vecY = genAllocCopy(rewriter, loc, memY, tokens);
   genBlockingWait(rewriter, loc, tokens);
   tokens.clear();
 
   // Create sparse environment and sparse matrix/dense vector handles.
   Type indexTp = rewriter.getIndexType();
-  Type handleTp = rewriter.getType<gpu::SparseHandleType>();
+  Type envHandleTp = rewriter.getType<gpu::SparseEnvHandleType>();
+  Type dnVecHandleTp = rewriter.getType<gpu::SparseDnVecHandleType>();
+  Type spmatHandleTp = rewriter.getType<gpu::SparseSpMatHandleType>();
   Type tokenTp = rewriter.getType<gpu::AsyncTokenType>();
   Value token = genFirstWait(rewriter, loc);
   auto env =
-      rewriter.create<gpu::CreateSparseEnvOp>(loc, handleTp, tokenTp, token);
+      rewriter.create<gpu::CreateSparseEnvOp>(loc, envHandleTp, tokenTp, token);
   Value handle = env.getResult(0);
   token = env.getAsyncToken();
-  Operation *spGenA = genSpMat(rewriter, loc, handleTp, tokenTp, token, szY,
-                               szX, nseA, rowA, colA, valA, isCOO, enableRT);
+  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, handleTp, tokenTp,
+  auto dvecX = rewriter.create<gpu::CreateDnVecOp>(loc, dnVecHandleTp, tokenTp,
                                                    token, vecX, szX);
   Value dnX = dvecX.getResult(0);
   token = dvecX.getAsyncToken();
-  auto dvecY = rewriter.create<gpu::CreateDnVecOp>(loc, handleTp, tokenTp,
+  auto dvecY = rewriter.create<gpu::CreateDnVecOp>(loc, dnVecHandleTp, tokenTp,
                                                    token, vecY, szY);
   Value dnY = dvecY.getResult(0);
   token = dvecY.getAsyncToken();
 
   // Precompute buffersize for SpMV.
   auto bufferComp = rewriter.create<gpu::SpMVBufferSizeOp>(
       loc, indexTp, tokenTp, token, handle, spMatA, dnX, dnY);
   Value bufferSz = bufferComp.getResult(0);
static LogicalResult rewriteSpMM(PatternRewriter &rewriter,
   Value matB = genAllocCopy(rewriter, loc, bufB, tokens);
   Value bufC = genTensorToMemref(rewriter, loc, c);
   Value matC = genAllocCopy(rewriter, loc, bufC, tokens);
   genBlockingWait(rewriter, loc, tokens);
   tokens.clear();
 
   // Create sparse environment and sparse matrix/dense matrix handles.
   Type indexTp = rewriter.getIndexType();
-  Type handleTp = rewriter.getType<gpu::SparseHandleType>();
+  Type envHandleTp = rewriter.getType<gpu::SparseEnvHandleType>();
+  Type dnMatHandleTp = rewriter.getType<gpu::SparseDnMatHandleType>();
+  Type spMatHandleTp = rewriter.getType<gpu::SparseSpMatHandleType>();
   Type tokenTp = rewriter.getType<gpu::AsyncTokenType>();
   Value token = genFirstWait(rewriter, loc);
   auto env =
-      rewriter.create<gpu::CreateSparseEnvOp>(loc, handleTp, tokenTp, token);
+      rewriter.create<gpu::CreateSparseEnvOp>(loc, envHandleTp, tokenTp, token);
   Value handle = env.getResult(0);
   token = env.getAsyncToken();
-  Operation *spGenA = genSpMat(rewriter, loc, handleTp, tokenTp, token, szm,
-                               szk, nseA, rowA, colA, valA, isCOO, enableRT);
+  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, handleTp, tokenTp,
+  auto dmatB = rewriter.create<gpu::CreateDnMatOp>(loc, dnMatHandleTp, tokenTp,
                                                    token, szk, szn, matB);
   Value dnB = dmatB.getResult(0);
   token = dmatB.getAsyncToken();
-  auto dmatC = rewriter.create<gpu::CreateDnMatOp>(loc, handleTp, tokenTp,
+  auto dmatC = rewriter.create<gpu::CreateDnMatOp>(loc, dnMatHandleTp, tokenTp,
                                                    token, szm, szn, matC);
   Value dnC = dmatC.getResult(0);
   token = dmatC.getAsyncToken();
 
   // Precompute buffersize for SpMM.
   auto bufferComp = rewriter.create<gpu::SpMMBufferSizeOp>(
       loc, indexTp, tokenTp, token, handle, spMatA, dnB, dnC);
   Value bufferSz = bufferComp.getResult(0);