[mlir][sparse][gpu] recognizing sddmm pattern in GPU libgen path

[K-Wu]

rebase origin/main with sddmm op added

[K-Wu]

clang-format

[K-Wu]

rebase origin/main

[K-Wu]

add test

[K-Wu]

involving the test

[K-Wu]

mlir/test/Dialect/SparseTensor/GPU/gpu_sampled_matmul_lib.mlir

88

can we instead output #CSR type here?

[K-Wu]

incorporating Aart's test case

[K-Wu]

add an in-place (from above, i.e., the program input) op test

[K-Wu]

format

[K-Wu]

saving progress

[K-Wu]

fixing compile error; remove repeated test case

[K-Wu]

adding more tests

[K-Wu]

update comment

[K-Wu]

rebase origin/main

[K-Wu]

drafting rewrite rule

[K-Wu]

saving progress

[K-Wu]

add rules and update tests

[K-Wu]

fix filechck parsing error

[K-Wu]

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

943

rewrite rule here

[K-Wu]

mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sampled-matmul-lib.mlir

43

TODO: add checking logic

[K-Wu]

mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sampled-matmul-lib.mlir

43

Didn't notice there is already checking logic down below. Closing this

[K-Wu]

mlir/test/Dialect/SparseTensor/GPU/gpu_sampled_matmul_lib.mlir

98

%1 = sparse_tensor_in_place_result %args: tensor<8x8xf64, #SM> // something new, I have to think about it

%2 = linalg.generic #trait_sampled_dense_dense
  ins(%args, %arga, %argb: tensor<8x8xf64, #SM>,
                           tensor<8x8xf64>, tensor<8x8xf64>)
  outs(%1: tensor<8x8xf64, #SM>) {
    ^bb(%s: f64, %a: f64, %b: f64, %x: f64):
      %p = arith.mulf %a, %b : f64
      %q = arith.mulf %s, %p : f64
      %r = arith.addf %x, %q : f64
      linalg.yield %r : f64
} -> tensor<8x8xf64, #SM>
return %2 : tensor<8x8xf64, #SM>

[K-Wu]

fix integration test

[K-Wu]

rebase; clean up a bit rule

[K-Wu]

reformat

[aartbik]

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

669

we test that C is sparse, but no test at all for D

So, the code below will sometimes crash, since it uses information on C not D
(e.g. isCOO)

674

this comment says more than is done below, is this a note to self, or something else?

678

you miss the d : row here

690

This assumes that D has exactly the same pattern already as C (or more precisely, the right length for all the arrays into which we copy the result). Perhaps we should start with simply returning C again after this operation, and find a way to allow this at IR level.

769

This is the wrong replacement, we need

rewriter.replaceOpWithNewOp<sparse_tensor::LoadOp>(op, ....)

and probably without inserts.

mlir/test/Dialect/SparseTensor/GPU/gpu_sampled_matmul_lib.mlir

111

can you remove this part with reuse completely, I think we will just need the pattern above for proper testing

mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sampled-matmul-lib.mlir

47

dense output, does not go through the GPU kernel then, right?

[K-Wu]

upd test cases to reflect actual need though not passing yet

[K-Wu]

incorporating aart's ad hoc solution

[aartbik]

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

663

please fix 80-col breaks here in comment

665

I think we need a

if (op.getDpsInitOperand(0) != c)
return failure();

here to enforce the ad-hoc pattern

mlir/test/Dialect/SparseTensor/GPU/gpu_sampled_matmul_lib.mlir

105

perhaps add

TODO: find a cleaner way to express this

mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sampled-matmul-lib.mlir

48

Put the comment about the ad-hoc idiom here too, so we don't forget.
Also the TODO.

93

Just curious if you really want to use the from-file set up of the sampling matrix? We can also just create this in memory like so many other examples, right?

[aartbik]

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

665

Or probably more precisely after preproc

%0 = bufferization.alloc_tensor() copy(%arg0)

op.getDpsInitOperand(0).getDefiningOp

links to

%0 = bufferization.alloc_tensor() copy(%c)

769

Make sure to return

sparse_tensor.load %arg0

instead of just %arg0, since the rematerialization needs to be made explicit.
You will find an example of this in e.g. the method sparse_simply_dynamic1() in the sparse_out.mlir test

[K-Wu]

still fixing

[K-Wu]

addressing comments

[K-Wu]

correcting the sddmm rewriter and tests except for the recognition rule

[K-Wu]

addressing comment

[K-Wu]

incorporating the new pattern

[K-Wu]

passed tests it seems

[K-Wu]

mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sampled-matmul-lib.mlir

93

Okay I will incorporate it next time

[Peiming]

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

334–351

The nest is too deep here. Could you either try to break it into more functions or early returns on some simple conditions?

e.g.,

if (!isa_and_nonnull(def)) 
   return false;

Value s_out = xxx
xxx

[aartbik]

mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sampled-matmul-lib.mlir

48

We can now use https://reviews.llvm.org/D152969
and remove all comments about hacks, since this is proper in-place
definition that matches cuSPARSE precisely.

[aartbik]

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

274

actually with the new formulation, you no longer need this since this is proper

S(I,J) = SPY( S(I,J) ) x SUM_K A(I,K) B(K,J) * alpha + beta * S(I,J)

[K-Wu]

removing obsolete comments

[K-Wu]

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

334–351

Good point and example! Incorporated

[K-Wu]

incorporating comments

[Peiming]

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

335

You need to do a null check too

[K-Wu]

addressing comment

[K-Wu]

rebase origin/main

[aartbik]

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

332

Add comment

// Helper to detect .....

344

return s_out == use->getOperand(0); ?

348

This seems way too simple. You pretty much match any

%1 = sparse_tensor.unary
%2 = sparse_tensor.reduce

but you don't check that %1 produces A * B and %2 produces sum accumulation

It will match the current example, but also much more! Please fix

687

remove this TODO. The current idiom has in-place semantics that will be correctly dealt with by bufferization

mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sampled-matmul-lib.mlir

116

releasing %0 feels cleaner, so that we have the in-place behavior clear

[K-Wu]

rebase

[K-Wu]

addressing comments

[K-Wu]

fixing compile errors

[aartbik]

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

335

You will often find the idiom

if (Operation *def = ....getDefiningOp()) {

use def

}

for this purpose, which you can even combine with the isa

if (Operation *def getDefiningOp<...>

353

technically, you have not checked if unary feeds into reduce, right?
or did I miss that?

[K-Wu]

fix errors

[K-Wu]

fixing runtime error

[K-Wu]

fix format

[K-Wu]

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

344

This is checked two times in the loop so I don't find an easy way to simplify the logic

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

if (isa<arith::AddFOp>(def) || isa<arith::AddIOp>(def)) {
               matchMulOfArgs(op, def->getOperand(1))) ||
              (def->getOperand(1) == x &&
               matchMulOfArgs(op, def->getOperand(0)));
     }
   }
   return false;
 }
 
+static bool matchReductionOfUnary(linalg::GenericOp op) {
+  auto yieldOp = cast<linalg::YieldOp>(op.getRegion().front().getTerminator());
+  if (auto *def = yieldOp.getOperand(0).getDefiningOp()) {
+    if (isa<sparse_tensor::ReduceOp>(def)) {
+      Value s_out = op.getBlock()->getArguments()[2];
+      for (auto *use : s_out.getUsers()) {
+        if (isa<sparse_tensor::ReduceOp>(use) ||
+            isa<sparse_tensor::UnaryOp>(use)) {
+          if (s_out != use->getOperand(0)) {
+            // The sparse matrix should be specified as the pattern in the two
+            // operators.
+            return false;
+          }
+        } else {
+          return false;
+        }
+      }
+    }
+    return true;
+  } else {
+    return false;
+  }
+}
+
 /// Test for sorted COO with suitable data and coordinates types.
 static bool isAdmissibleCOO(SparseTensorType &aTp) {
   return aTp.isCompressedLvl(0) && aTp.isOrderedLvl(0) && !aTp.isUniqueLvl(0) &&
          aTp.isSingletonLvl(1) && aTp.isOrderedLvl(1) && aTp.isUniqueLvl(1) &&
          (aTp.getElementType().isF64() || aTp.getElementType().isF32()) &&
          (aTp.getCrdWidth() == 0 || aTp.getCrdWidth() == 32 ||
           aTp.getCrdWidth() == 64);
 }
static LogicalResult rewriteSpMM(PatternRewriter &rewriter,
   genBlockingWait(rewriter, loc, tokens);
   tokens.clear();
 
   // Done.
   rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, bufC);
   return success();
 }
 
+/// Match and rewrite SDDMM kernel.
+static LogicalResult rewriteSDDMM(PatternRewriter &rewriter,
+                                  linalg::GenericOp op, bool enableRT) {
+
+  // For now, this pass reuses C and copies the result non-zero elements to
+  // overwrite C's.
+  // As an ad hoc solution, this pass also assumes the linalg takes a,b,c as
+  // input argument, and c as the output. It recognizes this pattern and rewrite
+  // it.
+
+  Location loc = op.getLoc();
+  Value a = op.getOperand(1);
+  Value b = op.getOperand(2);
+  Value c = op.getOperand(0);
+
+  SmallVector<Value> tokens;
+
+  // Only admissible sparse matrix format and dense matrices.
+  bool isCOO = false;
+  SparseTensorType aTp = getSparseTensorType(a);
+  SparseTensorType bTp = getSparseTensorType(b);
+  SparseTensorType cTp = getSparseTensorType(c);
+
+  if (!areAdmissibleTypes(cTp, bTp, aTp, enableRT, false, isCOO))
+    return failure();
+
+  // cusparse currently does not support COO in its SDDMM kernel.
+  if (isCOO) {
+    return failure();
+  }
+
+  // The SDDMM does the in-place operation.
+  // TODO: If the sparse matrix C is reused, e.g., later on when we remove
+  // redundant memcpy, we may need to duplicate it before the operation so that
+  // users could use the new copy instead. Start sparse kernel and copy data
+  // from host to device.
+  //   a : bufA           -> matA
+  //   b : bufB           -> matA
+  //   c : memR/memC/memV -> rowC,colC,valC
+  Value nseC = rewriter.create<NumberOfEntriesOp>(loc, c);
+  Value szm = linalg::createOrFoldDimOp(rewriter, loc, a, 0);
+  Value szk = linalg::createOrFoldDimOp(rewriter, loc, a, 1);
+  Value szn = linalg::createOrFoldDimOp(rewriter, loc, b, 1);
+  Value bufA = genTensorToMemref(rewriter, loc, a);
+  Value matA = genAllocCopy(rewriter, loc, bufA, tokens);
+  Value bufB = genTensorToMemref(rewriter, loc, b);
+  Value matB = genAllocCopy(rewriter, loc, bufB, tokens);
+  Value memR = genFirstPosOrCrds(rewriter, loc, c, isCOO, enableRT);
+  Value memC = genSecondCrds(rewriter, loc, c, isCOO, enableRT);
+  Value memV = genToValues(rewriter, loc, c);
+  Value rowC = genAllocCopy(rewriter, loc, memR, tokens);
+  Value colC = memC ? genAllocCopy(rewriter, loc, memC, tokens) : Value();
+  Value valC = genAllocCopy(rewriter, loc, memV, tokens);
+  genBlockingWait(rewriter, loc, tokens);
+  tokens.clear();
+
+  // Create sparse environment and sparse matrix/dense matrix handles.
+  Type indexTp = rewriter.getIndexType();
+  Type envHandleTp = rewriter.getType<gpu::SparseEnvHandleType>();
+  Type dnMatHandleTp = rewriter.getType<gpu::SparseDnTensorHandleType>();
+  Type spMatHandleTp = rewriter.getType<gpu::SparseSpMatHandleType>();
+  Type tokenTp = rewriter.getType<gpu::AsyncTokenType>();
+  Value token = genFirstWait(rewriter, loc);
+  auto env =
+      rewriter.create<gpu::CreateSparseEnvOp>(loc, envHandleTp, tokenTp, token);
+  Value handle = env.getResult(0);
+  token = env.getAsyncToken();
+
+  auto dmatA = rewriter.create<gpu::CreateDnTensorOp>(
+      loc, dnMatHandleTp, tokenTp, token, handle, matA,
+      SmallVector<Value>{szm, szk});
+  Value dnA = dmatA.getResult(0);
+  token = dmatA.getAsyncToken();
+  auto dmatB = rewriter.create<gpu::CreateDnTensorOp>(
+      loc, dnMatHandleTp, tokenTp, token, handle, matB,
+      SmallVector<Value>{szk, szn});
+  Value dnB = dmatB.getResult(0);
+  token = dmatB.getAsyncToken();
+
+  Operation *spGenC =
+      genSpMat(rewriter, loc, spMatHandleTp, tokenTp, token, szm, szn, nseC,
+               rowC, colC, valC, isCOO, enableRT);
+  Value spMatC = spGenC->getResult(0);
+  token = spGenC->getResult(1);
+
+  auto dnCType = llvm::cast<ShapedType>(c.getType()).getElementType();
+  // Precompute buffersize for SDDMM.
+  auto bufferComp = rewriter.create<gpu::SDDMMBufferSizeOp>(
+      loc, indexTp, tokenTp, token, handle, dnA, dnB, spMatC, dnCType);
+  Value bufferSz = bufferComp.getResult(0);
+  token = bufferComp.getAsyncToken();
+  auto buf = genAllocBuffer(rewriter, loc, bufferSz, token);
+  Value buffer = buf.getResult(0);
+  token = buf.getAsyncToken();
+
+  // Perform the SDDMM.
+  auto sddmmComp = rewriter.create<gpu::SDDMMOp>(
+      loc, tokenTp, token, handle, dnA, dnB, spMatC, dnCType, buffer);
+  token = sddmmComp.getAsyncToken();
+
+  // Copy data back to host and free all the resoures.
+  token = rewriter.create<gpu::DestroyDnTensorOp>(loc, tokenTp, token, dnA)
+              .getAsyncToken();
+  token = rewriter.create<gpu::DestroyDnTensorOp>(loc, tokenTp, token, dnB)
+              .getAsyncToken();
+  token = rewriter.create<gpu::DestroySpMatOp>(loc, tokenTp, token, spMatC)
+              .getAsyncToken();
+  token = rewriter.create<gpu::DestroySparseEnvOp>(loc, tokenTp, token, handle)
+              .getAsyncToken();
+  token = genDeallocMemRef(rewriter, loc, buffer, token);
+  token = genDeallocMemRef(rewriter, loc, matA, token);
+  token = genDeallocMemRef(rewriter, loc, matB, token);
+  token = genDeallocMemRef(rewriter, loc, rowC, token);
+  if (colC)
+    token = genDeallocMemRef(rewriter, loc, colC, token);
+  token = genCopyMemRef(rewriter, loc, memV, valC, token);
+  token = genDeallocMemRef(rewriter, loc, valC, token);
+  tokens.push_back(token);
+  genBlockingWait(rewriter, loc, tokens);
+  tokens.clear();
+
+  rewriter.replaceOpWithNewOp<sparse_tensor::LoadOp>(op, op.getOperand(0));
+
+  // Erasing the dead and illegal alloc_tensor operation applied to the first
+  // input argument to the linalg operator.
+  // We need to make sure the first input argument is linked to
+  // bufferization::alloc_tensor(%args0)
+  if (op.getDpsInitOperand(0)
+          ->get()
+          .getDefiningOp<bufferization::AllocTensorOp>() != nullptr) {
+    rewriter.eraseOp(op.getDpsInitOperand(0)
+                         ->get()
+                         .getDefiningOp<bufferization::AllocTensorOp>());
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Rewriting rules for direct code generation.
 //===----------------------------------------------------------------------===//
 
 /// Proof-of-concept rewriter. This rule generates a GPU implementation
 /// for each outermost forall loop generated by the sparse compiler.
 /// TODO: right works with parallelization-strategy=dense-outer-loop
 ///       but give this its own flags in the future
if (numLoops == 3 && numTensors == 3 &&
         linalg::isParallelIterator(iteratorTypes[1]) &&
         linalg::isReductionIterator(iteratorTypes[2]) &&
         // TODO: add transposed {i, k}, {k, j}
         // TODO: maybe add transposed {i, j} in future
         maps == infer({{i, k}, {k, j}, {i, j}}) && matchSumOfMultOfArgs(op)) {
       return rewriteSpMM(rewriter, op, enableRT);
     }
 
+    // Recognize a SDDMM kernel.
+    if (numLoops == 3 && numTensors == 3 &&
+        linalg::isParallelIterator(iteratorTypes[0]) &&
+        linalg::isParallelIterator(iteratorTypes[1]) &&
+        linalg::isReductionIterator(iteratorTypes[2]) &&
+        // TODO: add transposed {i, k}, {k, j}
+        // TODO: maybe add transposed {i, j} in future
+        maps == infer({{i, k}, {k, j}, {i, j}}) && matchReductionOfUnary(op)) {
+      return rewriteSDDMM(rewriter, op, enableRT);
+    }
+
     return failure();
   }
 
 private:
   bool enableRT;
 };
 
 } // namespace

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

#define ALPHABETA(dtp, alpha, beta)                                            \
   } else if (dtp == CUDA_R_32F || dtp == CUDA_C_32F) {                         \
     (alpha##p) = reinterpret_cast<void *>(&(alpha##f));                        \
     (beta##p) = reinterpret_cast<void *>(&(beta##f));                          \
   } else {                                                                     \
     (alpha##p) = reinterpret_cast<void *>(&(alpha##d));                        \
     (beta##p) = reinterpret_cast<void *>(&(beta##d));                          \
   }
 
+#define ZERO(dtp, zero)                                                        \
+  __nv_bfloat16(zero##16bf) = 0.0f;                                            \
+  __half(zero##16f) = 0.0f;                                                    \
+  float(zero##f) = 0.0f;                                                       \
+  double(zero##d) = 0.0;                                                       \
+  const void *(zero##p) = nullptr;                                             \
+  if (dtp == CUDA_R_16BF || dtp == CUDA_C_16BF) {                              \
+    (zero##p) = reinterpret_cast<void *>(&(zero##16bf));                       \
+  } else if (dtp == CUDA_R_16F || dtp == CUDA_C_16F) {                         \
+    (zero##p) = reinterpret_cast<void *>(&(zero##16f));                        \
+  } else if (dtp == CUDA_R_32F || dtp == CUDA_C_32F) {                         \
+    (zero##p) = reinterpret_cast<void *>(&(zero##f));                          \
+  } else {                                                                     \
+    (zero##p) = reinterpret_cast<void *>(&(zero##d));                          \
+  }
+
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void *
 mgpuCreateSparseEnv(CUstream /*stream*/) {
   cusparseHandle_t handle = nullptr;
   CUSPARSE_REPORT_IF_ERROR(cusparseCreate(&handle))
   return reinterpret_cast<void *>(handle);
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
mgpuSDDMMBufferSize(void *h, int32_t ma, int32_t mb, void *a, void *b, void *c,
   cusparseHandle_t handle = reinterpret_cast<cusparseHandle_t>(h);
   cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma);
   cusparseOperation_t modeB = static_cast<cusparseOperation_t>(mb);
   cusparseDnMatDescr_t matA = reinterpret_cast<cusparseDnMatDescr_t>(a);
   cusparseDnMatDescr_t matB = reinterpret_cast<cusparseDnMatDescr_t>(b);
   cusparseSpMatDescr_t matC = reinterpret_cast<cusparseSpMatDescr_t>(c);
   auto cTp = static_cast<cudaDataType_t>(ctp);
   ALPHABETA(cTp, alpha, beta)
+  ZERO(cTp, zero)
   size_t bufferSize = 0;
   CUSPARSE_REPORT_IF_ERROR(cusparseSDDMM_bufferSize(
-      handle, modeA, modeB, alphap, matA, matB, betap, matC, cTp,
+      handle, modeA, modeB, alphap, matA, matB, zerop, matC, cTp,
       CUSPARSE_SDDMM_ALG_DEFAULT, &bufferSize))
   return bufferSize == 0 ? 1 : bufferSize; // avoid zero-alloc
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
 mgpuSDDMM(void *h, int32_t ma, int32_t mb, void *a, void *b, void *c,
           int32_t ctp, void *buf, CUstream /*stream*/) {
   cusparseHandle_t handle = reinterpret_cast<cusparseHandle_t>(h);
   cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma);
   cusparseOperation_t modeB = static_cast<cusparseOperation_t>(mb);
   cusparseDnMatDescr_t matA = reinterpret_cast<cusparseDnMatDescr_t>(a);
   cusparseDnMatDescr_t matB = reinterpret_cast<cusparseDnMatDescr_t>(b);
   cusparseSpMatDescr_t matC = reinterpret_cast<cusparseSpMatDescr_t>(c);
   auto cTp = static_cast<cudaDataType_t>(ctp);
   ALPHABETA(cTp, alpha, beta)
+  ZERO(cTp, zero)
   CUSPARSE_REPORT_IF_ERROR(cusparseSDDMM(handle, modeA, modeB, alphap, matA,
-                                         matB, betap, matC, cTp,
+                                         matB, zerop, matC, cTp,
                                          CUSPARSE_SDDMM_ALG_DEFAULT, buf))
 }
 
 #ifdef MLIR_ENABLE_CUDA_CUSPARSELT
 
 ///
 /// Wrapper methods for the cuSparseLt library.
 ///

mlir/test/Dialect/SparseTensor/GPU/gpu_sampled_matmul_lib.mlir

+// RUN: mlir-opt %s --linalg-generalize-named-ops \
+// RUN:             --sparsification="enable-gpu-libgen" | FileCheck %s
+
+#trait_sampled_dense_dense = {
+  indexing_maps = [
+  affine_map<(i,j,k) -> (i,k)>,  // A
+  affine_map<(i,j,k) -> (k,j)>,  // B
+  affine_map<(i,j,k) -> (i,j)>   // S (out)
+  ],
+  iterator_types = ["parallel", "parallel", "reduction"],
+  doc = "X(i,j) += S(i,j) SUM_k A(i,k) B(k,j)"
+}
+
+#trait_vec_op = {
+  indexing_maps = [
+  affine_map<(i,j) -> (i,j)>,  // a (in)
+  affine_map<(i,j) -> (i,j)>,  // b (in)
+  affine_map<(i,j) -> (i,j)>   // x (out)
+  ],
+  iterator_types = ["parallel", "parallel"]
+}
+
+#CSR = #sparse_tensor.encoding<{ lvlTypes = [ "dense", "compressed" ] }>
+
+module {
+
+// CHECK-LABEL:   func.func @sparse_sampled_dd(
+// CHECK-SAME:                                 %[[VAL_0:.*]]: tensor<8x8xf64, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "compressed" ] }>>,
+// CHECK-SAME:                                 %[[VAL_1:.*]]: tensor<8x8xf64>,
+// CHECK-SAME:                                 %[[VAL_2:.*]]: tensor<8x8xf64>) -> tensor<8x8xf64, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "compressed" ] }>> {
+// CHECK:           %[[VAL_3:.*]] = arith.constant 8 : index
+// CHECK:           %[[VAL_4:.*]] = arith.constant 0 : index
+// CHECK:           %[[VAL_5:.*]] = sparse_tensor.number_of_entries %[[VAL_0]] : tensor<8x8xf64, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "compressed" ] }>>
+// CHECK:           %[[VAL_6:.*]] = bufferization.to_memref %[[VAL_1]] : memref<8x8xf64>
+// CHECK:           %[[VAL_7:.*]] = gpu.wait async
+// CHECK:           %[[VAL_8:.*]], %[[VAL_9:.*]] = gpu.alloc async {{\[}}%[[VAL_7]]] () : memref<8x8xf64>
+// CHECK:           %[[VAL_10:.*]] = gpu.memcpy async {{\[}}%[[VAL_9]]] %[[VAL_8]], %[[VAL_6]] : memref<8x8xf64>, memref<8x8xf64>
+// CHECK:           %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : memref<8x8xf64>
+// CHECK:           %[[VAL_12:.*]] = gpu.wait async
+// CHECK:           %[[VAL_13:.*]], %[[VAL_14:.*]] = gpu.alloc async {{\[}}%[[VAL_12]]] () : memref<8x8xf64>
+// CHECK:           %[[VAL_15:.*]] = gpu.memcpy async {{\[}}%[[VAL_14]]] %[[VAL_13]], %[[VAL_11]] : memref<8x8xf64>, memref<8x8xf64>
+// CHECK:           %[[VAL_16:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 1 : index} : tensor<8x8xf64, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "compressed" ] }>> to memref<?xindex>
+// CHECK:           %[[VAL_17:.*]] = sparse_tensor.coordinates %[[VAL_0]] {level = 1 : index} : tensor<8x8xf64, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "compressed" ] }>> to memref<?xindex>
+// CHECK:           %[[VAL_18:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<8x8xf64, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "compressed" ] }>> to memref<?xf64>
+// CHECK:           %[[VAL_19:.*]] = gpu.wait async
+// CHECK:           %[[VAL_20:.*]] = memref.dim %[[VAL_16]], %[[VAL_4]] : memref<?xindex>
+// CHECK:           %[[VAL_21:.*]], %[[VAL_22:.*]] = gpu.alloc async {{\[}}%[[VAL_19]]] (%[[VAL_20]]) : memref<?xindex>
+// CHECK:           %[[VAL_23:.*]] = gpu.memcpy async {{\[}}%[[VAL_22]]] %[[VAL_21]], %[[VAL_16]] : memref<?xindex>, memref<?xindex>
+// CHECK:           %[[VAL_24:.*]] = gpu.wait async
+// CHECK:           %[[VAL_25:.*]] = memref.dim %[[VAL_17]], %[[VAL_4]] : memref<?xindex>
+// CHECK:           %[[VAL_26:.*]], %[[VAL_27:.*]] = gpu.alloc async {{\[}}%[[VAL_24]]] (%[[VAL_25]]) : memref<?xindex>
+// CHECK:           %[[VAL_28:.*]] = gpu.memcpy async {{\[}}%[[VAL_27]]] %[[VAL_26]], %[[VAL_17]] : memref<?xindex>, memref<?xindex>
+// CHECK:           %[[VAL_29:.*]] = gpu.wait async
+// CHECK:           %[[VAL_30:.*]] = memref.dim %[[VAL_18]], %[[VAL_4]] : memref<?xf64>
+// CHECK:           %[[VAL_31:.*]], %[[VAL_32:.*]] = gpu.alloc async {{\[}}%[[VAL_29]]] (%[[VAL_30]]) : memref<?xf64>
+// CHECK:           %[[VAL_33:.*]] = gpu.memcpy async {{\[}}%[[VAL_32]]] %[[VAL_31]], %[[VAL_18]] : memref<?xf64>, memref<?xf64>
+// CHECK:           gpu.wait {{\[}}%[[VAL_10]], %[[VAL_15]], %[[VAL_23]], %[[VAL_28]], %[[VAL_33]]]
+// CHECK:           %[[VAL_34:.*]] = gpu.wait async
+// CHECK:           %[[VAL_35:.*]], %[[VAL_36:.*]] = gpu.create_sparse_env async {{\[}}%[[VAL_34]]]
+// CHECK:           %[[VAL_37:.*]], %[[VAL_38:.*]] = gpu.create_dn_tensor async {{\[}}%[[VAL_36]]] %[[VAL_35]], %[[VAL_8]], %[[VAL_3]], %[[VAL_3]] : index, index into memref<8x8xf64>
+// CHECK:           %[[VAL_39:.*]], %[[VAL_40:.*]] = gpu.create_dn_tensor async {{\[}}%[[VAL_38]]] %[[VAL_35]], %[[VAL_13]], %[[VAL_3]], %[[VAL_3]] : index, index into memref<8x8xf64>
+// CHECK:           %[[VAL_41:.*]], %[[VAL_42:.*]] = gpu.create_csr async {{\[}}%[[VAL_40]]] %[[VAL_3]], %[[VAL_3]], %[[VAL_5]], %[[VAL_21]], %[[VAL_26]], %[[VAL_31]] : memref<?xindex>, memref<?xindex>, memref<?xf64>
+// CHECK:           %[[VAL_43:.*]], %[[VAL_44:.*]] = gpu.sddmm_buffer_size async {{\[}}%[[VAL_42]]] %[[VAL_35]], %[[VAL_37]], %[[VAL_39]], %[[VAL_41]] into f64
+// CHECK:           %[[VAL_45:.*]], %[[VAL_46:.*]] = gpu.alloc async {{\[}}%[[VAL_44]]] (%[[VAL_43]]) : memref<?xi8>
+// CHECK:           %[[VAL_47:.*]] = gpu.sddmm async {{\[}}%[[VAL_46]]] %[[VAL_35]], %[[VAL_37]], %[[VAL_39]], %[[VAL_41]], %[[VAL_45]] : memref<?xi8> into f64
+// CHECK:           %[[VAL_48:.*]] = gpu.destroy_dn_tensor async {{\[}}%[[VAL_47]]] %[[VAL_37]]
+// CHECK:           %[[VAL_49:.*]] = gpu.destroy_dn_tensor async {{\[}}%[[VAL_48]]] %[[VAL_39]]
+// CHECK:           %[[VAL_50:.*]] = gpu.destroy_sp_mat async {{\[}}%[[VAL_49]]] %[[VAL_41]]
+// CHECK:           %[[VAL_51:.*]] = gpu.destroy_sparse_env async {{\[}}%[[VAL_50]]] %[[VAL_35]]
+// CHECK:           %[[VAL_52:.*]] = gpu.dealloc async {{\[}}%[[VAL_51]]] %[[VAL_45]] : memref<?xi8>
+// CHECK:           %[[VAL_53:.*]] = gpu.dealloc async {{\[}}%[[VAL_52]]] %[[VAL_8]] : memref<8x8xf64>
+// CHECK:           %[[VAL_54:.*]] = gpu.dealloc async {{\[}}%[[VAL_53]]] %[[VAL_13]] : memref<8x8xf64>
+// CHECK:           %[[VAL_55:.*]] = gpu.dealloc async {{\[}}%[[VAL_54]]] %[[VAL_21]] : memref<?xindex>
+// CHECK:           %[[VAL_56:.*]] = gpu.dealloc async {{\[}}%[[VAL_55]]] %[[VAL_26]] : memref<?xindex>
+// CHECK:           %[[VAL_57:.*]] = gpu.memcpy async {{\[}}%[[VAL_56]]] %[[VAL_18]], %[[VAL_31]] : memref<?xf64>, memref<?xf64>
+// CHECK:           %[[VAL_58:.*]] = gpu.dealloc async {{\[}}%[[VAL_57]]] %[[VAL_31]] : memref<?xf64>
+// CHECK:           gpu.wait {{\[}}%[[VAL_58]]]
+// CHECK:           %[[VAL_59:.*]] = sparse_tensor.load %[[VAL_0]] : tensor<8x8xf64, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "compressed" ] }>>
+// CHECK:           return %[[VAL_59]] : tensor<8x8xf64, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "compressed" ] }>>
+// CHECK:         }
+//
+// A kernel that computes a direct sampled matrix matrix multiplication
+// (with sparse result).
+// Compute SDDMM C = C\spy AB
+// 
+func.func @sparse_sampled_dd(%argS: tensor<8x8xf64, #CSR>,
+                               %argA: tensor<8x8xf64>,
+                               %argB: tensor<8x8xf64>) -> tensor<8x8xf64, #CSR> {
+    // TODO: find a cleaner way to express this
+    // We currently use this slightly strange idiom with the same sparse
+    // matrix as input and output to tell the GPU code generator this operation
+    // can be done "in place" on argS, and the caller will only free the
+    // resources once. In the long run, we need better ways to express
+    // this in the IR
+    %result = linalg.generic #trait_sampled_dense_dense
+      ins(%argA, %argB: tensor<8x8xf64>, tensor<8x8xf64>)
+      outs(%argS: tensor<8x8xf64, #CSR>) {
+        ^bb(%a: f64, %b: f64, %s: f64):
+           %f0 = arith.constant 0.0 : f64
+           %u = sparse_tensor.unary %s : f64 to f64
+             present={
+                ^bb0(%p: f64):
+                  %mul = arith.mulf %a, %b : f64
+                  sparse_tensor.yield %mul : f64
+             }
+             absent={}
+           %r = sparse_tensor.reduce %s, %u, %f0 : f64 {
+              ^bb0(%p: f64, %q: f64):
+                %add = arith.addf %p, %q : f64
+                sparse_tensor.yield %add : f64
+            }
+           linalg.yield %r : f64
+    } -> tensor<8x8xf64, #CSR>
+    return %result : tensor<8x8xf64, #CSR>
+  }
+
+}

mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sampled-matmul-lib.mlir

+//
+// NOTE: this test requires gpu-sm80
+//
+// RUN: mlir-opt %s \
+// RUN:   --sparse-compiler="enable-runtime-library=true enable-gpu-libgen gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71"  \
+// RUN: | TENSOR0="%mlir_src_dir/test/Integration/data/test.mtx" \
+// RUN:   mlir-cpu-runner \
+// RUN:   --shared-libs=%mlir_cuda_runtime \
+// RUN:   --shared-libs=%mlir_c_runner_utils \
+// RUN:   --e entry --entry-point-result=void \
+// RUN: | FileCheck %s
+//
+
+!Filename = !llvm.ptr<i8>
+
+#CSR = #sparse_tensor.encoding<{
+  lvlTypes = ["dense", "compressed"]
+}>
+
+#trait_sampled_dense_dense = {
+  indexing_maps = [
+    affine_map<(i,j,k) -> (i,k)>,  // A
+    affine_map<(i,j,k) -> (k,j)>,  // B
+    affine_map<(i,j,k) -> (i,j)>   // S (out)
+  ],
+  iterator_types = ["parallel", "parallel", "reduction"],
+  doc = "X(i,j) += S(i,j) SUM_k A(i,k) B(k,j)"
+}
+
+//
+// Integration test that lowers a kernel annotated as sparse to
+// actual sparse code, initializes a matching sparse storage scheme
+// from file, and runs the resulting code with the JIT compiler.
+//
+module {
+  //
+  // A kernel that computes a sampled matrix matrix multiplication.
+  //
+  func.func @sampled_dense_dense(%args: tensor<?x?xf32, #CSR>,
+                                 %arga: tensor<?x?xf32>,
+                                 %argb: tensor<?x?xf32>) -> tensor<?x?xf32, #CSR> {
+    // TODO: find a cleaner way to express this
+    // We currently use this slightly strange idiom with the same sparse
+    // matrix as input and output to tell the GPU code generator this operation
+    // can be done "in place" on argS, and the caller will only free the
+    // resources once. In the long run, we need better ways to express
+    // this in the IR
+    %result = linalg.generic #trait_sampled_dense_dense
+      ins(%arga, %argb: tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%args: tensor<?x?xf32, #CSR>) {
+        ^bb(%a: f64, %b: f64, %s: f64):
+           %f0 = arith.constant 0.0 : f64
+           %u = sparse_tensor.unary %s : f64 to f64
+             present={
+                ^bb0(%p: f64):
+                  %mul = arith.mulf %a, %b : f64
+                  sparse_tensor.yield %mul : f64
+             }
+             absent={}
+           %r = sparse_tensor.reduce %s, %u, %f0 : f64 {
+              ^bb0(%p: f64, %q: f64):
+                %add = arith.addf %p, %q : f64
+                sparse_tensor.yield %add : f64
+            }
+           linalg.yield %r : f64
+      } -> tensor<?x?xf32, #CSR>
+    return %result : tensor<?x?xf32, #CSR>
+  }
+
+  func.func private @getTensorFilename(index) -> (!Filename)
+
+  //
+  // Main driver that reads matrix from file and calls the sparse kernel.
+  //
+  func.func @entry() {
+    %d0 = arith.constant 0.0 : f32
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %c5 = arith.constant 5 : index
+    %c10 = arith.constant 10 : index
+
+    // Initialize dense matrices.
+    %x = tensor.generate %c5, %c5 {
+    ^bb0(%i : index, %j : index):
+      tensor.yield %d0 : f32
+    } : tensor<?x?xf32>
+
+    %a = tensor.generate %c5, %c10 {
+    ^bb0(%i: index, %j: index):
+      %p = arith.addi %i, %c1 : index
+      %q = arith.index_cast %p : index to i32
+      %d = arith.sitofp %q : i32 to f32
+      tensor.yield %d : f32
+    } : tensor<?x?xf32>
+
+    %b = tensor.generate %c10, %c5 {
+    ^bb0(%i: index, %j: index):
+      %p = arith.addi %j, %c1 : index
+      %q = arith.index_cast %p : index to i32
+      %d = arith.sitofp %q : i32 to f32
+      tensor.yield %d : f32
+    } : tensor<?x?xf32>
+
+    // Read the sparse matrix from file, construct sparse storage.
+    %fileName = call @getTensorFilename(%c0) : (index) -> (!Filename)
+    %s = sparse_tensor.new %fileName : !Filename to tensor<?x?xf32, #CSR>
+
+    // Call the kernel.
+    %0 = call @sampled_dense_dense(%s, %a, %b)
+       : (tensor<?x?xf32, #CSR>,
+          tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32, #CSR>
+
+    // Print the result for verification.
+    //
+    // CHECK: ( 10, 40, 40, 100, 90, 40, 160, 100, 250 )
+    %vm = sparse_tensor.values %0 : tensor<?x?xf32, #CSR> to memref<?xf32>
+    %vv = vector.transfer_read %vm[%c0], %d0 : memref<?xf32>, vector<9xf32>
+    vector.print %vv : vector<9xf32>
+
+    // Release the resources.
+    bufferization.dealloc_tensor %s : tensor<?x?xf32, #CSR>
+
+    return
+  }
+}