[mlir][gpu][sparse] adding initial cusparse sddmm libgen support

[K-Wu]

Summary

This diff added the sddmm cuSparse op + roundtrip + lowering to wrapper lib

[K-Wu]

impl SDDMM representation and rewriter

[K-Wu]

add test cases

[aartbik]

I would split this into adding the new cuSparse op + roundtrip + lowering to wrapper lib in this revision
and then the actual recognition during sparsification + test + integration test perhaps in one or two next revisions

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

695

Note that I made a small refinement in https://reviews.llvm.org/D151404 which seems safer: keep them all on the same stream (so no blocking wait here) and move the copy after the copy-in only and buffer are released

[K-Wu]

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

695

noted. I will update this and mark done once I incorporate the rebased-pull into this diff.

[aartbik]

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

699

sparse output will be more complex than this btw

[K-Wu]

rebase origin main

[K-Wu]

fixing compile errors

[K-Wu]

fixing test error and formatting

[K-Wu]

rebase origin/main

[K-Wu]

fix compile error

[K-Wu]

fix format

[K-Wu]

rebase origin/main

[K-Wu]

removing SparseGPUCodegen.cpp from this diff list

[Peiming]

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

2025

bad indentation here

[K-Wu]

fix inden

[K-Wu]

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

2025

Good catch. Thank you!

[aartbik]

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

1944

matrix -> matrices (since we have dense A,B now)

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1490

note that we have this test + convert test so many times now, that perhaps a helper will be useful
(next revision is fine)

[K-Wu]

fix grammar

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

def GPU_SpMMBufferSizeOp : GPU_Op<"spmm_buffer_size", [GPU_AsyncOpInterface]> {
   let results = (outs Res<Index>:$bufferSz, Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $env `,` $spmatA `,` $dnmatB `,` $dnmatC attr-dict
   }];
 }
 
+// TODO: add GPU_SDDMMOp
+
 def GPU_SpMMOp : GPU_Op<"spmm", [GPU_AsyncOpInterface]> {
   let summary = "SpMM operation";
   let description = [{
     The `gpu.spmm` operation performs the SpMM operation on the given sparse and
     dense matrix, and buffer.  The operation expects handles returned by previous
     sparse operations to construct an environment and the operands for SpMM. The
     buffer must have been allocated on the device.
 
def GPU_SpMMOp : GPU_Op<"spmm", [GPU_AsyncOpInterface]> {
   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

public:
   }
 
 private:
   LogicalResult
   matchAndRewrite(gpu::SpMMBufferSizeOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+// TODO: impl SDDMM Op lowering pass here
+
 class ConvertSpMMOpToGpuRuntimeCallPattern
     : public ConvertOpToGpuRuntimeCallPattern<gpu::SpMMOp> {
 public:
   ConvertSpMMOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter)
       : ConvertOpToGpuRuntimeCallPattern<gpu::SpMMOp>(typeConverter) {}
 
 private:
   LogicalResult
void mlir::populateGpuToLLVMConversionPatterns(LLVMTypeConverter &converter,
   });
   converter.addConversion([&converter](gpu::SparseHandleType type) -> Type {
     return converter.getPointerType(
         IntegerType::get(&converter.getContext(), 8));
   });
 
   patterns.add<ConvertAllocOpToGpuRuntimeCallPattern,
                ConvertDeallocOpToGpuRuntimeCallPattern,
                ConvertHostRegisterOpToGpuRuntimeCallPattern,
                ConvertHostUnregisterOpToGpuRuntimeCallPattern,
                ConvertMemcpyOpToGpuRuntimeCallPattern,
                ConvertMemsetOpToGpuRuntimeCallPattern,
                ConvertSetDefaultDeviceOpToGpuRuntimeCallPattern,
                ConvertWaitAsyncOpToGpuRuntimeCallPattern,
                ConvertWaitOpToGpuRuntimeCallPattern,
                ConvertAsyncYieldToGpuRuntimeCallPattern,
                ConvertCreateSparseEnvOpToGpuRuntimeCallPattern,

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

static LogicalResult rewriteSpMV(PatternRewriter &rewriter,
   genBlockingWait(rewriter, loc, tokens);
   tokens.clear();
 
   // Done.
   rewriter.replaceOp(op, op.getDpsInitOperand(0)->get());
   return success();
 }
 
+// TODO: implement SDDMM rewriter here
+
 /// Match and rewrite SpMM kernel.
 static LogicalResult rewriteSpMM(PatternRewriter &rewriter,
                                  linalg::GenericOp op, bool enableRT) {
   Location loc = op.getLoc();
   Value a = op.getOperand(0);
   Value b = op.getOperand(1);
   Value c = op.getOperand(2); // we have C = AB
   SmallVector<Value> tokens;
LogicalResult matchAndRewrite(scf::ParallelOp forallOp,
     tokens.clear();
     Value kernelToken =
         genLaunchGPUFunc(rewriter, gpuFunc, args, tokens, numThreads);
     // Finalize the outlined arguments.
     genParametersOut(rewriter, loc, out, kernelToken, scalars, buffers, args,
                      tokens);
     genBlockingWait(rewriter, loc, tokens);
     rewriter.eraseOp(forallOp);
     return success();
   }
 
 private:
   // Helper method to see if block appears in given loop.
   static bool isNestedIn(Block *block, scf::ParallelOp forallOp) {
     for (Operation *o = block->getParentOp(); o; o = o->getParentOp()) {
       if (o == forallOp)
         return true;
     }
     return false;
   }