[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

166–173

move this up to other includes

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

169

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

198

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

166–173

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

169

That makes sense. I remove that

198

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

174–175

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

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

1481–1497

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

174–175

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

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

1560

Worked on it

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1481–1497

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

195

this is nice!

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

150

same order as in decl, so spmat last

209

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

219–220

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

225

here too

[K-Wu]

addressing the order of refined HandleType

[K-Wu]

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

195

credit to

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

219–220

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

195

Great thanks to @wrengr and @Peiming

[K-Wu]

centralize keywords

[K-Wu]

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

219–220

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

219–220

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

1699

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>()">, "sparse dense vector handle type">,
+             BuildableType<"mlir::gpu::SparseDnVecHandleType::get($_builder.getContext())">;
+def GPU_SparseDnMatHandle : DialectType<
+  GPU_Dialect, CPred<"$_self.isa<::mlir::gpu::SparseDnMatHandleType>()">, "sparse 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 in sparse">,
+             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.
 
     GPU operations implementing this interface take a list of dependencies
     as `gpu.async.token` arguments and optionally return a `gpu.async.token`.
 
     The op doesn't start executing until all depent ops producing the async
     dependency tokens have finished executing.
 
     If the op returns a token, the op merely schedules the execution on the
     device and returns immediately, without waiting for the execution to
     complete. On the hand, if the op does not return a token, the op will wait

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.
-class SparseHandleType
-    : public Type::TypeBase<SparseHandleType, Type, TypeStorage> {
+#include "llvm/Support/Casting.h"
+
+// Handle types for sparse.
+enum class SparseHandleKind { Env, DnVec, DnMat, SpMat, Default };
+
+template <SparseHandleKind K>
+class SparseHandleConcreteType
+    : public Type::TypeBase<SparseHandleConcreteType<K>, Type, TypeStorage> {
 public:
-  // Used for generic hooks in TypeBase.
-  using Base::Base;
+  using Type::TypeBase<SparseHandleConcreteType<K>, Type,
+                       TypeStorage>::Base::Base;
+  SparseHandleKind getKind() const { return kind; }
+
+  SparseHandleConcreteType(SparseHandleKind kind) : kind(kind) {}
+  SparseHandleConcreteType() = default;
+  template <SparseHandleKind OtherK>
+  SparseHandleConcreteType(SparseHandleConcreteType<OtherK> other)
+      : SparseHandleConcreteType(K) {}
+  static bool classof(mlir::Type S) {
+    if (!llvm::isa<SparseHandleConcreteType>(S))
+      return false;
+    auto SS = llvm::dyn_cast<SparseHandleConcreteType>(S);
+    return SS.getKind() == K;
+  }
+
+private:
+  const SparseHandleKind kind = SparseHandleKind::Default;
 };
 
+using SparseDnVecHandleType = SparseHandleConcreteType<SparseHandleKind::DnVec>;
+using SparseDnMatHandleType = SparseHandleConcreteType<SparseHandleKind::DnMat>;
+using SparseSpMatHandleType = SparseHandleConcreteType<SparseHandleKind::SpMat>;
+using SparseEnvHandleType = SparseHandleConcreteType<SparseHandleKind::Env>;
+
 } // 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"

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]> {
   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
   }];
 }
 
 def GPU_CreateDnVecOp : GPU_Op<"create_dn_vec", [GPU_AsyncOpInterface]> {
let description = [{
 
     ```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
   }];
 }
 
 def GPU_CreateDnMatOp : GPU_Op<"create_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);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) $dmat attr-dict
   }];
 }
 
 def GPU_CreateCooOp : GPU_Op<"create_coo", [GPU_AsyncOpInterface]> {
def GPU_CreateCooOp : GPU_Op<"create_coo", [GPU_AsyncOpInterface]> {
 
   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);
+                   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
   }];
 }
 
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)
   }];
 }

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

 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 {
+  converter.addConversion(
+      [&converter](gpu::SparseDnVecHandleType type) -> Type {
+        return converter.getPointerType(
+            IntegerType::get(&converter.getContext(), 8));
+      });
+  converter.addConversion(
+      [&converter](gpu::SparseDnMatHandleType type) -> Type {
+        return converter.getPointerType(
+            IntegerType::get(&converter.getContext(), 8));
+      });
+  converter.addConversion(
+      [&converter](gpu::SparseSpMatHandleType type) -> Type {
+        return converter.getPointerType(
+            IntegerType::get(&converter.getContext(), 8));
+      });
+  converter.addConversion([&converter](gpu::SparseEnvHandleType type) -> Type {
     return converter.getPointerType(
         IntegerType::get(&converter.getContext(), 8));
   });
 
   patterns.add<ConvertAllocOpToGpuRuntimeCallPattern,
                ConvertDeallocOpToGpuRuntimeCallPattern,
                ConvertHostRegisterOpToGpuRuntimeCallPattern,
                ConvertHostUnregisterOpToGpuRuntimeCallPattern,
                ConvertMemcpyOpToGpuRuntimeCallPattern,

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<SparseSpMatHandleType>();
+  addTypes<SparseDnMatHandleType>();
+  addTypes<SparseDnVecHandleType>();
   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"
       >();
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 == "sparse.env_handle")
+    return SparseEnvHandleType::get(context);
+  if (keyword == "sparse.spmat_handle")
+    return SparseSpMatHandleType::get(context);
+  if (keyword == "sparse.dnvec_handle")
+    return SparseDnVecHandleType::get(context);
+  if (keyword == "sparse.dnmat_handle")
+    return SparseDnMatHandleType::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 << "sparse.env_handle"; })
+      .Case<SparseSpMatHandleType>([&](Type) { os << "sparse.spmat_handle"; })
+      .Case<SparseDnMatHandleType>([&](Type) { os << "sparse.dnmat_handle"; })
+      .Case<SparseDnVecHandleType>([&](Type) { os << "sparse.dnvec_handle"; })
       .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, nnzA,
+               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,
   auto env =
       rewriter.create<gpu::CreateSparseEnvOp>(loc, handleTp, 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);
   Value spMatA = spGenA->getResult(0);
   token = spGenA->getResult(1);
   auto dmatB = rewriter.create<gpu::CreateDnMatOp>(loc, handleTp, tokenTp,
                                                    token, szk, szn, matB);
   Value dnB = dmatB.getResult(0);
   token = dmatB.getAsyncToken();
   auto dmatC = rewriter.create<gpu::CreateDnMatOp>(loc, handleTp, tokenTp,
                                                    token, szm, szn, matC);
   Value dnC = dmatC.getResult(0);
   token = dmatC.getAsyncToken();