[mlir][sparse][gpu] only support default algorithm for SpGEMM

[aartbik]

Summary

Rationale:
This is the approach taken for all the others too (SpMV, SpMM, SDDMM),
so it is more consistent to follow the same path (until we have a need
for more algorithms). Also, in a follow up revision, this will allow
us to remove some unused GEMM ops.

[aartbik]

fixed edit

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

def GPU_SDDMMOp : GPU_Op<"sddmm", [GPU_AsyncOpInterface]> {
   ];
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $dnmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $spmatC `,` $buffer attr-dict `:` type($buffer) `into` $computeType
   }];
 }
 
-// ALG1, ALG2, ALG3 use 3--5 to align with cusparseSpGEMMAlg_t in cusparse.h.
-def GPU_SpGEMMAlg : I32EnumAttr<"SpGEMMAlg",
-    "selected algorithm for sparse matrix SpGEMM",
-    [
-      I32EnumAttrCase<"ALG1", 3>,
-      I32EnumAttrCase<"ALG2", 4>,
-      I32EnumAttrCase<"ALG3", 5>,
-    ]> {
-      let genSpecializedAttr = 0;
-      let cppNamespace = GPU_Dialect.cppNamespace;
-      let defaultValue = "SpGEMMAlg::ALG1";
-}
-
-def GPU_SpGEMMAlgAttr : EnumAttr<GPU_Dialect, GPU_SpGEMMAlg, "spgemm_alg"> {
-  let defaultValue = GPU_SpGEMMAlg.defaultValue;
-}
-
 def GPU_SpGEMMWorkEstimationOrComputeKind : I32EnumAttr<"SpGEMMWorkEstimationOrComputeKind",
     "choose whether spgemm_work_estimation_or_compute does work estimation or compute",
     [
       I32EnumAttrCase<"WORK_ESTIMATION", 0>,
       I32EnumAttrCase<"COMPUTE", 1>,
     ]> {
       let genSpecializedAttr = 0;
       let cppNamespace = GPU_Dialect.cppNamespace;
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
-    %desc,  %token = gpu.spgemm_create_descr async [%dep]
+    %desc, %token = gpu.spgemm_create_descr async [%dep]
     ```
-
   }];
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies);
   let results = (outs GPU_SparseSpGEMMOpHandle:$desc,
                       Optional<GPU_AsyncToken>:$asyncToken);
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     attr-dict
   }];
 }
 
 def GPU_SpGEMMDestroyDescrOp : GPU_Op<"spgemm_destroy_descr", [GPU_AsyncOpInterface]> {
   let summary = "SpGEMM Destroy Descr operation";
   let description = [{
     The `gpu.spgemm_destroy_descr` destroys the SpGEMM operation 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
     %token = gpu.spgemm_destroy_descr async [%dep] %desc
     ```
-
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                        GPU_SparseSpGEMMOpHandle:$desc);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $desc attr-dict
   }];
 }
 
-
 def GPU_SpGEMMWorkEstimationOrComputeOp : GPU_Op<"spgemm_work_estimation_or_compute", [GPU_AsyncOpInterface]> {
   let summary = "SpGEMM work estimation operation";
   let description = [{
     The `gpu.spgemm_work_estimation_or_compute` is used to call
     cusparseSpGEMM_workEstimation or cusparseSpGEMM_compute. Both of them are
     for both determining the buffer size and performing the actual computation.
     The operation expects handles returned by previous sparse operations to
     construct an environment and the operands for SpGEMM.
     The buffer must have been allocated on the device.
 
-
     C' = alpha * op(A) * op(B) + beta * C
 
     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
     %bufferSz, %token = gpu.spgemm_work_estimation_or_compute async [%dep]{COMPUTE}
                           %desc, %spmatA{NON_TRANSPOSE}, %spmatB{NON_TRANSPOSE},
                           %spmatC, ALG2, %spgemmDesc, %c0, %alloc: f32 into
                           memref<0xi8>
     ```
 
     The matrix arguments can also be associated with one of the following
     operators: NON_TRANSPOSE, TRANSPOSE, CONJUGATE_TRANSPOSE. The default value
     is NON_TRANSPOSE.
-
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                        GPU_SparseSpGEMMOpHandle:$desc,
                        GPU_TransposeModeAttr:$modeA,
                        GPU_TransposeModeAttr:$modeB,
                        GPU_SparseSpMatHandle:$spmatA,
                        GPU_SparseSpMatHandle:$spmatB,
                        GPU_SparseSpMatHandle:$spmatC,
                        TypeAttr:$computeType,
                        Index:$bufferSz,
-                       GPU_SpGEMMAlgAttr:$alg,
                        AnyMemRef:$buffer,
                        GPU_SpGEMMWorkEstimationOrComputeKindAttr:$kind);
   let results = (outs Res<Index>:$bufferSzNew,
                       Optional<GPU_AsyncToken>:$asyncToken);
 
   let builders = [OpBuilder<(ins
     "Type":$bufferSzNew,
     "Type":$asyncToken,
     "ValueRange":$asyncDependencies,
     "Value":$desc,
     "Value":$spmatA,
     "Value":$spmatB,
     "Value":$spmatC,
     "Type":$computeType,
     "Value":$bufferSz,
     "Value":$buffer), [{
   auto modeA = gpu::TransposeMode::NON_TRANSPOSE;
   auto modeB = gpu::TransposeMode::NON_TRANSPOSE;
-  auto alg = gpu::SpGEMMAlg::ALG1;
   auto kind = gpu::SpGEMMWorkEstimationOrComputeKind::WORK_ESTIMATION;
   return build($_builder, $_state, bufferSzNew, asyncToken, asyncDependencies, desc,
-               modeA, modeB, spmatA, spmatB, spmatC, computeType, bufferSz, alg, buffer, kind);}]>
+               modeA, modeB, spmatA, spmatB, spmatC, computeType, bufferSz, buffer, kind);}]>
   ];
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    `{` $kind `}` $spmatA (`{` $modeA^ `}`)? `,` $spmatB (`{` $modeB^ `}`)? `,` $spmatC `,` $alg `,` $desc `,` $bufferSz `,` $buffer  attr-dict `:` $computeType `into` type($buffer)
+    `{` $kind `}` $spmatA (`{` $modeA^ `}`)? `,` $spmatB (`{` $modeB^ `}`)? `,` $spmatC `,` $desc `,` $bufferSz `,` $buffer  attr-dict `:` $computeType `into` type($buffer)
   }];
 }
 
-
 def GPU_SpGEMMEstimateMemoryOp : GPU_Op<"spgemm_estimate_memory", [GPU_AsyncOpInterface]> {
   let summary = "SpGEMM estimate memory operation";
   let description = [{
     The `gpu.spgemm_estimate_memory` is used for both determining the buffer
     size and performing the actual computation.
 
     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
     %bufferSz3, %dummy, %token = gpu.spgemm_estimate_memory async [%dep] %spmatA, %spmatB, %spmatC, ALG2, %spgemmDesc, %c0, %c0, %alloc: f32 into memref<0xi8>
     ```
-
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                        GPU_SparseSpGEMMOpHandle:$desc,
                        GPU_TransposeModeAttr:$modeA,
                        GPU_TransposeModeAttr:$modeB,
                        GPU_SparseSpMatHandle:$spmatA,
                        GPU_SparseSpMatHandle:$spmatB,
                        GPU_SparseSpMatHandle:$spmatC,
                        TypeAttr:$computeType,
-                       GPU_SpGEMMAlgAttr:$alg,
                        Index:$bufferSz3,
                        AnyMemRef:$buffer3,
                        Index:$bufferSz2);
   let results = (outs Index:$bufferSz3New,
                       Index:$bufferSz2New,
                       Optional<GPU_AsyncToken>:$asyncToken);
 
   let builders = [OpBuilder<(ins
     "Type":$bufferSz3New,
     "Type":$bufferSz2New,
     "Type":$asyncToken,
     "ValueRange":$asyncDependencies,
     "Value":$desc,
     "Value":$spmatA,
     "Value":$spmatB,
     "Value":$spmatC,
     "Type":$computeType,
     "Value":$bufferSz3,
     "Value":$buffer3,
     "Value":$bufferSz2), [{
   auto modeA = gpu::TransposeMode::NON_TRANSPOSE;
   auto modeB = gpu::TransposeMode::NON_TRANSPOSE;
-  auto alg = gpu::SpGEMMAlg::ALG1;
   return build($_builder, $_state, bufferSz3New, bufferSz2New, asyncToken,
                asyncDependencies, desc, modeA, modeB, spmatA, spmatB, spmatC,
-               computeType, alg, bufferSz3, buffer3, bufferSz2);}]>
+               computeType, bufferSz3, buffer3, bufferSz2);}]>
   ];
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $spmatA (`{` $modeA^ `}`)? `,` $spmatB (`{` $modeB^ `}`)? `,` $spmatC `,` $alg `,` $desc `,` $bufferSz3 `,` $bufferSz2 `,` $buffer3 attr-dict `:` $computeType `into` type($buffer3)
+    $spmatA (`{` $modeA^ `}`)? `,` $spmatB (`{` $modeB^ `}`)? `,` $spmatC `,` $desc `,` $bufferSz3 `,` $bufferSz2 `,` $buffer3 attr-dict `:` $computeType `into` type($buffer3)
   }];
 }
 
-
 def GPU_SpGEMMCopyOp : GPU_Op<"spgemm_copy", [GPU_AsyncOpInterface]> {
   let summary = "SpGEMM copy operation";
   let description = [{
     The `gpu.spgemm_copy` operation copies a sparse matrix, e.g., the result of
     the SpGEMM computation.
 
     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
     gpu.spgemm_copy %spmatA, %spmatB, %spmatC, ALG2, %spgemmDesc: f32
     ```
 
     The matrix arguments can also be associated with one of the following
     operators: NON_TRANSPOSE, TRANSPOSE, CONJUGATE_TRANSPOSE. The default value
     is NON_TRANSPOSE.
-
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                        GPU_SparseSpGEMMOpHandle:$desc,
                        GPU_TransposeModeAttr:$modeA,
                        GPU_TransposeModeAttr:$modeB,
                        GPU_SparseSpMatHandle:$spmatA,
                        GPU_SparseSpMatHandle:$spmatB,
                        GPU_SparseSpMatHandle:$spmatC,
-                       TypeAttr:$computeType,
-                       GPU_SpGEMMAlgAttr:$alg);
+                       TypeAttr:$computeType);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let builders = [OpBuilder<(ins
     "Type":$asyncToken,
     "ValueRange":$asyncDependencies,
     "Value":$desc,
     "Value":$spmatA,
     "Value":$spmatB,
     "Value":$spmatC,
     "Type":$computeType), [{
   auto modeA = gpu::TransposeMode::NON_TRANSPOSE;
   auto modeB = gpu::TransposeMode::NON_TRANSPOSE;
-  auto alg = gpu::SpGEMMAlg::ALG1;
   return build($_builder, $_state, asyncToken, asyncDependencies, desc,
-               modeA, modeB, spmatA, spmatB, spmatC, computeType, alg);}]>
+               modeA, modeB, spmatA, spmatB, spmatC, computeType);}]>
   ];
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $spmatA (`{` $modeA^ `}`)? `,` $spmatB (`{` $modeB^ `}`)? `,` $spmatC `,` $alg `,` $desc attr-dict `:` $computeType
+    $spmatA (`{` $modeA^ `}`)? `,` $spmatB (`{` $modeB^ `}`)? `,` $spmatC `,` $desc attr-dict `:` $computeType
   }];
 }
 
-
 def GPU_SpGEMMGetSizeOp : GPU_Op<"spgemm_get_size", [GPU_AsyncOpInterface]> {
   let summary = "SpGEMM get size operation";
   let description = [{
     The `gpu.spgemm_get_size` operation retrieves the number of rows, number of
     columns, and number of non-zero elements of a sparse matrix.
 
     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
     %rows, %cols, %nnz, %token = gpu.spgemm_get_size async [%dep] %spmatC
     ```
-
-    The matrix arguments can also be associated with one of the following
-    operators: NON_TRANSPOSE, TRANSPOSE, CONJUGATE_TRANSPOSE. The default value
-    is NON_TRANSPOSE.
-
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                        GPU_SparseSpMatHandle:$spmat);
   let results = (outs Index:$rows,
                       Index:$cols,
                       Index:$nnz,
                       Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $spmat attr-dict
   }];
 }
 
 #endif // GPU_OPS

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

FunctionCallBuilder createCuSparseLtSpMMBuilder = {
       llvmVoidType,
       {llvmPointerType, llvmPointerType, llvmPointerType, llvmPointerType,
        llvmPointerType, llvmPointerType, llvmPointerType /*void *stream*/}};
   FunctionCallBuilder createSpGEMMWorkEstimationBuilder = {
       "mgpuSpGEMMWorkEstimation",
       llvmIntPtrType,
       {llvmPointerType /*s*/, llvmInt32Type /*ma*/, llvmInt32Type /*mb*/,
        llvmPointerType /*a*/, llvmPointerType /*b*/, llvmPointerType /*c*/,
-       llvmInt32Type /*ctp*/, llvmInt32Type /*alg*/, llvmIntPtrType /*bs*/,
-       llvmPointerType /*buf*/, llvmPointerType /*void *stream*/}};
+       llvmInt32Type /*ctp*/, llvmIntPtrType /*bs*/, llvmPointerType /*buf*/,
+       llvmPointerType /*void *stream*/}};
   FunctionCallBuilder createSpGEMMEstimateMemoryBuilder = {
       "mgpuSpGEMMEstimateMemory",
       llvmVoidType,
       {llvmPointerType /*nbs3*/, llvmPointerType /*nbs2*/,
        llvmPointerType /*s*/, llvmInt32Type /*ma*/, llvmInt32Type /*mb*/,
        llvmPointerType /*a*/, llvmPointerType /*b*/, llvmPointerType /*c*/,
-       llvmInt32Type /*ctp*/, llvmInt32Type /*alg*/,
-       llvmFloat32Type /*chunk_fraction*/, llvmIntPtrType /*bs3*/,
-       llvmPointerType /*buf3*/, llvmIntPtrType /*bs2*/,
+       llvmInt32Type /*ctp*/, llvmFloat32Type /*chunk_fraction*/,
+       llvmIntPtrType /*bs3*/, llvmPointerType /*buf3*/, llvmIntPtrType /*bs2*/,
        llvmPointerType /*void *stream*/}};
   FunctionCallBuilder createSpGEMMComputeBuilder = {
       "mgpuSpGEMMCompute",
       llvmIntPtrType,
       {llvmPointerType /*s*/, llvmInt32Type /*ma*/, llvmInt32Type /*mb*/,
        llvmPointerType /*a*/, llvmPointerType /*b*/, llvmPointerType /*c*/,
-       llvmInt32Type /*ctp*/, llvmInt32Type /*alg*/, llvmIntPtrType /*bs*/,
-       llvmPointerType /*buf*/, llvmPointerType /*void *stream*/}};
+       llvmInt32Type /*ctp*/, llvmIntPtrType /*bs*/, llvmPointerType /*buf*/,
+       llvmPointerType /*void *stream*/}};
   FunctionCallBuilder createSpGEMMCopyBuilder = {
       "mgpuSpGEMMCopy",
       llvmVoidType,
       {llvmPointerType /*s*/, llvmInt32Type /*ma*/, llvmInt32Type /*mb*/,
        llvmPointerType /*a*/, llvmPointerType /*b*/, llvmPointerType /*c*/,
-       llvmInt32Type /*ctp*/, llvmInt32Type /*alg*/,
-       llvmPointerType /*void *stream*/}};
+       llvmInt32Type /*ctp*/, llvmPointerType /*void *stream*/}};
   FunctionCallBuilder createSpGEMMCreateDescrBuilder = {
       "mgpuSpGEMMCreateDescr",
       llvmPointerType,
       {llvmPointerType /*void *stream*/}};
   FunctionCallBuilder createSpGEMMDestroyDescrBuilder = {
       "mgpuSpGEMMDestoryDescr",
       llvmVoidType,
       {llvmPointerType /*s*/, llvmPointerType /*void *stream*/}};
ConvertSpGEMMWorkEstimationOrComputeOpToGpuRuntimeCallPattern::matchAndRewrite(
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
   auto computeType = genConstInt32From(
       rewriter, loc, getCuSparseDataTypeFrom(adaptor.getComputeType()));
   auto modeA = genConstInt32From(rewriter, loc, adaptor.getModeA());
   auto modeB = genConstInt32From(rewriter, loc, adaptor.getModeB());
-  auto alg = genConstInt32From(rewriter, loc, adaptor.getAlg());
   auto stream = adaptor.getAsyncDependencies().front();
 
   Value pBuf =
       MemRefDescriptor(adaptor.getBuffer()).allocatedPtr(rewriter, loc);
   if (!getTypeConverter()->useOpaquePointers())
     pBuf = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, pBuf);
 
   Value bufferSizeNew;
 
   if (adaptor.getKind() ==
       gpu::SpGEMMWorkEstimationOrComputeKind::WORK_ESTIMATION) {
     bufferSizeNew =
         createSpGEMMWorkEstimationBuilder
             .create(loc, rewriter,
                     {adaptor.getDesc(), modeA, modeB, adaptor.getSpmatA(),
-                     adaptor.getSpmatB(), adaptor.getSpmatC(), computeType, alg,
+                     adaptor.getSpmatB(), adaptor.getSpmatC(), computeType,
                      adaptor.getBufferSz(), pBuf, stream})
             .getResult();
   } else {
     bufferSizeNew =
         createSpGEMMComputeBuilder
             .create(loc, rewriter,
                     {adaptor.getDesc(), modeA, modeB, adaptor.getSpmatA(),
-                     adaptor.getSpmatB(), adaptor.getSpmatC(), computeType, alg,
+                     adaptor.getSpmatB(), adaptor.getSpmatC(), computeType,
                      adaptor.getBufferSz(), pBuf, stream})
             .getResult();
   }
   rewriter.replaceOp(op, {bufferSizeNew, stream});
   return success();
 }
 
 LogicalResult
 ConvertSpGEMMEstimateMemoryOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::SpGEMMEstimateMemoryOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
   auto computeType = genConstInt32From(
       rewriter, loc, getCuSparseDataTypeFrom(adaptor.getComputeType()));
-  auto alg = genConstInt32From(rewriter, loc, adaptor.getAlg());
   auto modeA = genConstInt32From(rewriter, loc, adaptor.getModeA());
   auto modeB = genConstInt32From(rewriter, loc, adaptor.getModeB());
   auto stream = adaptor.getAsyncDependencies().front();
   // TODO: support other chunk fraction
   Value chunkFraction = genConstFloat32From(rewriter, loc, 1.0);
   Value pBuf3 =
       MemRefDescriptor(adaptor.getBuffer3()).allocatedPtr(rewriter, loc);
   if (!getTypeConverter()->useOpaquePointers())
auto bufferSizePtr3 = rewriter.create<LLVM::GEPOp>(
       loc, llvmInt64PointerType, llvmInt64PointerType, bufferSize,
       ValueRange{rewriter.create<LLVM::ConstantOp>(loc, getIndexType(),
                                                    rewriter.getIndexAttr(1))});
 
   createSpGEMMEstimateMemoryBuilder.create(
       loc, rewriter,
       {bufferSizePtr3, bufferSizePtr2, adaptor.getDesc(), modeA, modeB,
        adaptor.getSpmatA(), adaptor.getSpmatB(), adaptor.getSpmatC(),
-       computeType, alg, chunkFraction, adaptor.getBufferSz3(), pBuf3,
+       computeType, chunkFraction, adaptor.getBufferSz3(), pBuf3,
        adaptor.getBufferSz2(), stream});
   auto bufferSize2 =
       rewriter.create<LLVM::LoadOp>(loc, llvmInt64Type, bufferSizePtr2);
   auto bufferSize3 =
       rewriter.create<LLVM::LoadOp>(loc, llvmInt64Type, bufferSizePtr3);
 
   rewriter.replaceOp(op, {bufferSize3, bufferSize2, stream});
   return success();
 }
 
 LogicalResult ConvertSpGEMMCopyOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::SpGEMMCopyOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
   auto computeType = genConstInt32From(
       rewriter, loc, getCuSparseDataTypeFrom(adaptor.getComputeType()));
   auto modeA = genConstInt32From(rewriter, loc, adaptor.getModeA());
   auto modeB = genConstInt32From(rewriter, loc, adaptor.getModeB());
-  auto alg = genConstInt32From(rewriter, loc, adaptor.getAlg());
   auto stream = adaptor.getAsyncDependencies().front();
-  createSpGEMMCopyBuilder.create(
-      loc, rewriter,
-      {adaptor.getDesc(), modeA, modeB, adaptor.getSpmatA(),
-       adaptor.getSpmatB(), adaptor.getSpmatC(), computeType, alg, stream});
+  createSpGEMMCopyBuilder.create(loc, rewriter,
+                                 {adaptor.getDesc(), modeA, modeB,
+                                  adaptor.getSpmatA(), adaptor.getSpmatB(),
+                                  adaptor.getSpmatC(), computeType, stream});
   rewriter.replaceOp(op, {stream});
   return success();
 }
 
 LogicalResult ConvertSpGEMMGetSizeOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::SpGEMMGetSizeOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSDDMM(int32_t ma, int32_t mb,
   ALPHABETA(cTp, alpha, beta)
   CUSPARSE_REPORT_IF_ERROR(cusparseSDDMM(cusparse_env, modeA, modeB, alphap,
                                          matA, matB, betap, matC, cTp,
                                          CUSPARSE_SDDMM_ALG_DEFAULT, buf))
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t mgpuSpGEMMWorkEstimation(
     void *s, int32_t ma, int32_t mb, void *a, void *b, void *c, int32_t ctp,
-    int32_t alg, intptr_t bs, void *buf, CUstream /*stream*/) {
+    intptr_t bs, void *buf, CUstream /*stream*/) {
   cusparseSpGEMMDescr_t spgemmDesc = reinterpret_cast<cusparseSpGEMMDescr_t>(s);
   cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma);
   cusparseOperation_t modeB = static_cast<cusparseOperation_t>(mb);
-  cusparseSpGEMMAlg_t algorithm = static_cast<cusparseSpGEMMAlg_t>(alg);
   cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a);
   cusparseSpMatDescr_t matB = reinterpret_cast<cusparseSpMatDescr_t>(b);
   cusparseSpMatDescr_t matC = reinterpret_cast<cusparseSpMatDescr_t>(c);
   auto cTp = static_cast<cudaDataType_t>(ctp);
   ALPHABETA(cTp, alpha, beta)
   size_t newBufferSize = bs;
 
   CUSPARSE_REPORT_IF_ERROR(cusparseSpGEMM_workEstimation(
       cusparse_env, modeA, modeB, alphap, matA, matB, betap, matC, cTp,
-      algorithm, spgemmDesc, &newBufferSize, buf))
+      CUSPARSE_SPGEMM_DEFAULT, spgemmDesc, &newBufferSize, buf))
   return newBufferSize == 0 ? 1 : newBufferSize; // avoid zero-alloc
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
 mgpuSpGEMMEstimateMemory(void *nbs3, void *nbs2, void *s, int32_t ma,
                          int32_t mb, void *a, void *b, void *c, int32_t ctp,
-                         int32_t alg, float chunk_fraction, intptr_t bs3,
-                         void *buf3, intptr_t bs2, CUstream /*stream*/) {
+                         float chunk_fraction, intptr_t bs3, void *buf3,
+                         intptr_t bs2, CUstream /*stream*/) {
   cusparseSpGEMMDescr_t spgemmDesc = reinterpret_cast<cusparseSpGEMMDescr_t>(s);
   cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma);
   cusparseOperation_t modeB = static_cast<cusparseOperation_t>(mb);
   cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a);
   cusparseSpMatDescr_t matB = reinterpret_cast<cusparseSpMatDescr_t>(b);
   cusparseSpMatDescr_t matC = reinterpret_cast<cusparseSpMatDescr_t>(c);
   auto cTp = static_cast<cudaDataType_t>(ctp);
   ALPHABETA(cTp, alpha, beta)
   size_t *newBufferSize2 = reinterpret_cast<size_t *>(nbs2);
   size_t *newBufferSize3 = reinterpret_cast<size_t *>(nbs3);
   *newBufferSize2 = bs2;
   *newBufferSize3 = bs3;
-  auto algorithm = static_cast<cusparseSpGEMMAlg_t>(alg);
 
   CUSPARSE_REPORT_IF_ERROR(cusparseSpGEMM_estimateMemory(
       cusparse_env, modeA, modeB, alphap, matA, matB, betap, matC, cTp,
-      algorithm, spgemmDesc, chunk_fraction, newBufferSize3, buf3,
+      CUSPARSE_SPGEMM_DEFAULT, spgemmDesc, chunk_fraction, newBufferSize3, buf3,
       newBufferSize2))
   // avoid zero-alloc
   if (*newBufferSize2 == 0) {
     *newBufferSize2 = 1;
   }
   if (*newBufferSize3 == 0) {
     *newBufferSize3 = 1;
   }
   return;
 }
 
-extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t mgpuSpGEMMCompute(
-    void *s, int32_t ma, int32_t mb, void *a, void *b, void *c, int32_t ctp,
-    int32_t alg, intptr_t bsz2, void *buf2, CUstream /*stream*/) {
+extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t
+mgpuSpGEMMCompute(void *s, int32_t ma, int32_t mb, void *a, void *b, void *c,
+                  int32_t ctp, intptr_t bsz2, void *buf2, CUstream /*stream*/) {
   cusparseSpGEMMDescr_t spgemmDesc = reinterpret_cast<cusparseSpGEMMDescr_t>(s);
   cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma);
   cusparseOperation_t modeB = static_cast<cusparseOperation_t>(mb);
-  cusparseSpGEMMAlg_t algorithm = static_cast<cusparseSpGEMMAlg_t>(alg);
   cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a);
   cusparseSpMatDescr_t matB = reinterpret_cast<cusparseSpMatDescr_t>(b);
   cusparseSpMatDescr_t matC = reinterpret_cast<cusparseSpMatDescr_t>(c);
   auto cTp = static_cast<cudaDataType_t>(ctp);
   ALPHABETA(cTp, alpha, beta)
   size_t newBufferSize2 = bsz2;
   CUSPARSE_REPORT_IF_ERROR(cusparseSpGEMM_compute(
       cusparse_env, modeA, modeB, alphap, matA, matB, betap, matC, cTp,
-      algorithm, spgemmDesc, &newBufferSize2, buf2))
+      CUSPARSE_SPGEMM_DEFAULT, spgemmDesc, &newBufferSize2, buf2))
   return newBufferSize2 == 0 ? 1 : newBufferSize2; // avoid zero-alloc
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
 mgpuSpGEMMCopy(void *s, int32_t ma, int32_t mb, void *a, void *b, void *c,
-               int32_t ctp, int32_t alg, CUstream /*stream*/) {
+               int32_t ctp, CUstream /*stream*/) {
   cusparseSpGEMMDescr_t spgemmDesc = reinterpret_cast<cusparseSpGEMMDescr_t>(s);
   cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma);
   cusparseOperation_t modeB = static_cast<cusparseOperation_t>(mb);
   cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a);
   cusparseSpMatDescr_t matB = reinterpret_cast<cusparseSpMatDescr_t>(b);
   cusparseSpMatDescr_t matC = reinterpret_cast<cusparseSpMatDescr_t>(c);
   auto cTp = static_cast<cudaDataType_t>(ctp);
-  auto algorithm = static_cast<cusparseSpGEMMAlg_t>(alg);
   ALPHABETA(cTp, alpha, beta)
 
-  CUSPARSE_REPORT_IF_ERROR(cusparseSpGEMM_copy(cusparse_env, modeA, modeB,
-                                               alphap, matA, matB, betap, matC,
-                                               cTp, algorithm, spgemmDesc))
+  CUSPARSE_REPORT_IF_ERROR(
+      cusparseSpGEMM_copy(cusparse_env, modeA, modeB, alphap, matA, matB, betap,
+                          matC, cTp, CUSPARSE_SPGEMM_DEFAULT, spgemmDesc))
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void *
 mgpuSpGEMMCreateDescr(CUstream /*stream*/) {
-  // cusparseSpGEMMDescr_t is a pointer type
   cusparseSpGEMMDescr_t spgemmDesc = nullptr;
   CUSPARSE_REPORT_IF_ERROR(cusparseSpGEMM_createDescr(&spgemmDesc))
   return reinterpret_cast<void *>(spgemmDesc);
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
 mgpuSpGEMMDestroyDescr(void *s, CUstream /*stream*/) {
-  // cusparseSpGEMMDescr_t is a pointer type
   cusparseSpGEMMDescr_t spgemmDesc = reinterpret_cast<cusparseSpGEMMDescr_t>(s);
   CUSPARSE_REPORT_IF_ERROR(cusparseSpGEMM_destroyDescr(spgemmDesc))
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
 mgpuSpGEMMGetSize(void *m, void *r, void *c, void *n, CUstream /*stream*/) {
   cusparseConstSpMatDescr_t matDescr =
       reinterpret_cast<cusparseConstSpMatDescr_t>(m);
CUSPARSE_REPORT_IF_ERROR(
                        /*dD*/ matC->values, d_workspace, nullptr, 0))
 
   CUSPARSE_REPORT_IF_ERROR(cusparseLtMatDescriptorDestroy(&(matA->mat)))
   // destroy the plan associated with the sparse matrix
   CUSPARSE_REPORT_IF_ERROR(cusparseLtMatmulPlanDestroy(&(matA->plan)))
 }
 
 #endif // MLIR_ENABLE_CUDA_CUSPARSELT
-#endif // MLIR_ENABLE_CUDA_CUSPARSE

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

func.func @spgemm(%arg0: index) {
     %mem2, %token2 = gpu.alloc async [%token1] (%arg0) : memref<?xf64>
     %spmatA, %token3 = gpu.create_csr async [%token2] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
     %spmatB, %token4 = gpu.create_csr async [%token3] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
     %spmatC, %token5 = gpu.create_csr async [%token4] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
     %spgemmDesc, %token6 = gpu.spgemm_create_descr async [%token5]
     // Used as nullptr
     %alloc = memref.alloc() : memref<0xi8>
     %c0 = arith.constant 0 : index
-    %bufferSz1, %token7 = gpu.spgemm_work_estimation_or_compute async 
+    %bufferSz1, %token7 = gpu.spgemm_work_estimation_or_compute async
                             [%token6]{WORK_ESTIMATION}
-                            %spmatA{NON_TRANSPOSE}, %spmatB{NON_TRANSPOSE}, 
-                            %spmatC, ALG2, %spgemmDesc, %c0, 
+                            %spmatA{NON_TRANSPOSE}, %spmatB{NON_TRANSPOSE},
+                            %spmatC, %spgemmDesc, %c0,
                             %alloc: f32 into memref<0xi8>
     %buf1, %token8 = gpu.alloc async [%token7] (%bufferSz1) : memref<?xi8>
     %bufferSz1_1, %token9 = gpu.spgemm_work_estimation_or_compute async
-                              [%token8]{WORK_ESTIMATION} %spmatA, %spmatB, 
-                              %spmatC, ALG2, %spgemmDesc, %bufferSz1, 
+                              [%token8]{WORK_ESTIMATION} %spmatA, %spmatB,
+                              %spmatC, %spgemmDesc, %bufferSz1,
                               %buf1: f32 into memref<?xi8>
-    %bufferSz3, %dummy, %token10 = gpu.spgemm_estimate_memory async [%token9] 
-                                     %spmatA, %spmatB, %spmatC, ALG2, 
-                                     %spgemmDesc, %c0, %c0, 
+    %bufferSz3, %dummy, %token10 = gpu.spgemm_estimate_memory async [%token9]
+                                     %spmatA, %spmatB, %spmatC,
+                                     %spgemmDesc, %c0, %c0,
                                      %alloc: f32 into memref<0xi8>
     %buf3, %token11 = gpu.alloc async [%token10] (%bufferSz3) : memref<?xi8>
-    %bufferSz3_2, %bufferSz2, %token12 = gpu.spgemm_estimate_memory async 
+    %bufferSz3_2, %bufferSz2, %token12 = gpu.spgemm_estimate_memory async
                                           [%token11] %spmatA, %spmatB, %spmatC,
-                                          ALG2, %spgemmDesc, %bufferSz3, %c0,
+                                          %spgemmDesc, %bufferSz3, %c0,
                                           %buf3: f32 into memref<?xi8>
     %buf2, %token13 = gpu.alloc async [%token12] (%bufferSz2) : memref<?xi8>
-    %bufferSz2_2, %token14 = gpu.spgemm_work_estimation_or_compute async 
-                               [%token13]{COMPUTE} %spmatA, %spmatB, %spmatC, 
-                               ALG2, %spgemmDesc, %bufferSz2, 
+    %bufferSz2_2, %token14 = gpu.spgemm_work_estimation_or_compute async
+                               [%token13]{COMPUTE} %spmatA, %spmatB, %spmatC,
+                               %spgemmDesc, %bufferSz2,
                                %buf2: f32 into memref<?xi8>
     %rows, %cols, %nnz, %token15 = gpu.spgemm_get_size async [%token14] %spmatC
     %mem_columns, %token16 = gpu.alloc async [%token15] (%cols) : memref<?xi32>
     %mem_values, %token17 = gpu.alloc async [%token16] (%nnz) : memref<?xf32>
     gpu.wait [%token17]
     %token18 = gpu.wait async
-    %token19 = gpu.spgemm_copy async [%token18] %spmatA, %spmatB, %spmatC, ALG2, %spgemmDesc: f32
+    %token19 = gpu.spgemm_copy async [%token18] %spmatA, %spmatB, %spmatC, %spgemmDesc: f32
     %token20 = gpu.destroy_sp_mat async [%token19] %spmatA
     %token21 = gpu.destroy_sp_mat async [%token20] %spmatB
     %token22 = gpu.destroy_sp_mat async [%token21] %spmatC
     gpu.wait [%token22]
     return
   }
 
 }
-
-

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

module attributes {gpu.container_module} {
   // CHECK:           %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xindex>
   // CHECK:           %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xf64>
   // CHECK:           %{{.*}}, %{{.*}} = gpu.create_csr async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xindex>, memref<?xindex>, memref<?xf64>
   // CHECK:           %{{.*}}, %{{.*}} = gpu.create_csr async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xindex>, memref<?xindex>, memref<?xf64>
   // CHECK:           %{{.*}}, %{{.*}} = gpu.create_csr async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xindex>, memref<?xindex>, memref<?xf64>
   // CHECK:           %{{.*}}, %{{.*}} = gpu.spgemm_create_descr async [%{{.*}}]
   // CHECK:           %{{.*}} = memref.alloc() : memref<0xi8>
   // CHECK:           %{{.*}} = arith.constant 0 : index
-  // CHECK:           %{{.*}}, %{{.*}} = gpu.spgemm_work_estimation_or_compute async [%{{.*}}]{{{.*}}} %{{.*}}, %{{.*}}, %{{.*}},  ALG2, %{{.*}}, %{{.*}}, %{{.*}} : f32 into memref<0xi8>
+  // CHECK:           %{{.*}}, %{{.*}} = gpu.spgemm_work_estimation_or_compute async [%{{.*}}]{{{.*}}} %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : f32 into memref<0xi8>
   // CHECK:           %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xi8>
-  // CHECK:           %{{.*}}, %{{.*}} = gpu.spgemm_work_estimation_or_compute async [%{{.*}}]{{{.*}}} %{{.*}}, %{{.*}}, %{{.*}},  ALG2, %{{.*}}, %{{.*}}, %{{.*}} : f32 into memref<?xi8>
-  // CHECK:           %{{.*}}, %{{.*}}, %{{.*}} = gpu.spgemm_estimate_memory async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}},  ALG2, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : f32 into memref<0xi8>
+  // CHECK:           %{{.*}}, %{{.*}} = gpu.spgemm_work_estimation_or_compute async [%{{.*}}]{{{.*}}} %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : f32 into memref<?xi8>
+  // CHECK:           %{{.*}}, %{{.*}}, %{{.*}} = gpu.spgemm_estimate_memory async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : f32 into memref<0xi8>
   // CHECK:           %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xi8>
-  // CHECK:           %{{.*}}, %{{.*}}, %{{.*}} = gpu.spgemm_estimate_memory async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}},  ALG2, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : f32 into memref<?xi8>
+  // CHECK:           %{{.*}}, %{{.*}}, %{{.*}} = gpu.spgemm_estimate_memory async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : f32 into memref<?xi8>
   // CHECK:           %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xi8>
-  // CHECK:           %{{.*}}, %{{.*}} = gpu.spgemm_work_estimation_or_compute async [%{{.*}}]{{{.*}}} %{{.*}}, %{{.*}}, %{{.*}},  ALG2, %{{.*}}, %{{.*}}, %{{.*}} : f32 into memref<?xi8>
+  // CHECK:           %{{.*}}, %{{.*}} = gpu.spgemm_work_estimation_or_compute async [%{{.*}}]{{{.*}}} %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : f32 into memref<?xi8>
   // CHECK:           %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} = gpu.spgemm_get_size async [%{{.*}}] %{{.*}}
   // CHECK:           %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xi32>
   // CHECK:           %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xf32>
   // CHECK:           gpu.wait [%{{.*}}]
-  // CHECK:           gpu.spgemm_copy  %{{.*}}, %{{.*}}, %{{.*}},  ALG2, %{{.*}} : f32
+  // CHECK:           gpu.spgemm_copy  %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : f32
   // CHECK:           gpu.destroy_sp_mat  %{{.*}}
   // CHECK:           gpu.destroy_sp_mat  %{{.*}}
   // CHECK:           gpu.destroy_sp_mat  %{{.*}}
   // CHECK:           return
   func.func @spgemm(%arg0: index) {
     %token0 = gpu.wait async
     %mem1, %token1 = gpu.alloc async [%token0] (%arg0) : memref<?xindex>
     %mem2, %token2 = gpu.alloc async [%token1] (%arg0) : memref<?xf64>
     %spmatA, %token3 = gpu.create_csr async [%token2] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
     %spmatB, %token4 = gpu.create_csr async [%token3] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
     %spmatC, %token5 = gpu.create_csr async [%token4] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
     %spgemmDesc, %token6 = gpu.spgemm_create_descr async [%token5]
     // Used as nullptr
     %alloc = memref.alloc() : memref<0xi8>
     %c0 = arith.constant 0 : index
-    %bufferSz1, %token7 = gpu.spgemm_work_estimation_or_compute async 
+    %bufferSz1, %token7 = gpu.spgemm_work_estimation_or_compute async
                             [%token6]{WORK_ESTIMATION}
-                            %spmatA{NON_TRANSPOSE}, %spmatB{NON_TRANSPOSE}, 
-                            %spmatC, ALG2, %spgemmDesc, %c0, 
+                            %spmatA{NON_TRANSPOSE}, %spmatB{NON_TRANSPOSE},
+                            %spmatC, %spgemmDesc, %c0,
                             %alloc: f32 into memref<0xi8>
     %buf1, %token8 = gpu.alloc async [%token7] (%bufferSz1) : memref<?xi8>
-    %bufferSz1_1, %token9 = gpu.spgemm_work_estimation_or_compute async 
-                              [%token8]{WORK_ESTIMATION} %spmatA, %spmatB, 
-                              %spmatC, ALG2, %spgemmDesc, %bufferSz1, 
+    %bufferSz1_1, %token9 = gpu.spgemm_work_estimation_or_compute async
+                              [%token8]{WORK_ESTIMATION} %spmatA, %spmatB,
+                              %spmatC, %spgemmDesc, %bufferSz1,
                               %buf1: f32 into memref<?xi8>
-    %bufferSz3, %dummy, %token10 = gpu.spgemm_estimate_memory async [%token9] 
-                                     %spmatA, %spmatB, %spmatC, ALG2, 
-                                     %spgemmDesc, %c0, %c0, 
+    %bufferSz3, %dummy, %token10 = gpu.spgemm_estimate_memory async [%token9]
+                                     %spmatA, %spmatB, %spmatC,
+                                     %spgemmDesc, %c0, %c0,
                                      %alloc: f32 into memref<0xi8>
     %buf3, %token11 = gpu.alloc async [%token10] (%bufferSz3) : memref<?xi8>
-    %bufferSz3_2, %bufferSz2, %token12 = gpu.spgemm_estimate_memory async 
+    %bufferSz3_2, %bufferSz2, %token12 = gpu.spgemm_estimate_memory async
                                           [%token11] %spmatA, %spmatB, %spmatC,
-                                          ALG2, %spgemmDesc, %bufferSz3, %c0,
+                                          %spgemmDesc, %bufferSz3, %c0,
                                           %buf3: f32 into memref<?xi8>
     %buf2, %token13 = gpu.alloc async [%token12] (%bufferSz2) : memref<?xi8>
-    %bufferSz2_2, %token14 = gpu.spgemm_work_estimation_or_compute async 
-                               [%token13]{COMPUTE} %spmatA, %spmatB, %spmatC, 
-                               ALG2, %spgemmDesc, %bufferSz2, 
+    %bufferSz2_2, %token14 = gpu.spgemm_work_estimation_or_compute async
+                               [%token13]{COMPUTE} %spmatA, %spmatB, %spmatC,
+                               %spgemmDesc, %bufferSz2,
                                %buf2: f32 into memref<?xi8>
     %rows, %cols, %nnz, %token15 = gpu.spgemm_get_size async [%token14] %spmatC
     %mem_columns, %token16 = gpu.alloc async [%token15] (%cols) : memref<?xi32>
     %mem_values, %token17 = gpu.alloc async [%token16] (%nnz) : memref<?xf32>
     gpu.wait [%token17]
-    gpu.spgemm_copy %spmatA, %spmatB, %spmatC, ALG2, %spgemmDesc: f32
+    gpu.spgemm_copy %spmatA, %spmatB, %spmatC, %spgemmDesc: f32
     gpu.destroy_sp_mat %spmatA
     gpu.destroy_sp_mat %spmatB
     gpu.destroy_sp_mat %spmatC
     return
   }
 
   // CHECK-LABEL: func @sddmm
   // CHECK: %{{.*}} = gpu.wait async
func.func @sddmm(%arg0: index) {
     %token7 = gpu.sddmm async [%token6] %dnmat, %dnmat, %spmat, %mem2 : memref<?xf64> into f64
     %token8 = gpu.destroy_sp_mat async [%token7] %spmat
     %token9 = gpu.destroy_dn_tensor async [%token8] %dnmat
     gpu.wait [%token9]
     return
   }
 
 }
-
-