[Mlir-commits] [mlir] [mlir][sparse][gpu] Add libgen support for transposed SpMV and SpMM (PR #230412)
Vito Secona
llvmlistbot at llvm.org
Fri Oct 9 01:47:33 PDT 2026
https://github.com/secona created https://github.com/llvm/llvm-project/pull/230412
This change adds libgen support for transposed SpMV and SpMM in the sparse-gpu-codegen pass by recognizing `y = A^Tx` for SpMV and `y = A^TB`, `y = AB^T`, and `y = A^TB^T` for SpMM. Tests have been added to the SparseTensor dialect for these changes. Tests are also added in the GPU dialect for lowering to GPU runtime calls as it was previously missing.
Assisted by Codex
>From ecd66fecbaf2b6e1c63547001d2cb1f0e6ce3e20 Mon Sep 17 00:00:00 2001
From: Vito Secona <secona00 at gmail.com>
Date: Fri, 9 Oct 2026 11:34:33 +0700
Subject: [PATCH 1/4] [mlir][gpu] Test transposed SpMV roundtrip and runtime
lowering
---
.../lower-sparse-to-gpu-runtime-calls.mlir | 34 +++++++++++++++++--
mlir/test/Dialect/GPU/sparse-roundtrip.mlir | 26 ++++++++++++++
2 files changed, 58 insertions(+), 2 deletions(-)
diff --git a/mlir/test/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir b/mlir/test/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir
index d4c0a76088356f..2e45813dc1e1c6 100644
--- a/mlir/test/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir
+++ b/mlir/test/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir
@@ -8,8 +8,10 @@ module attributes {gpu.container_module} {
// CHECK: llvm.call @mgpuMemAlloc
// CHECK: llvm.call @mgpuCreateCoo
// CHECK: llvm.call @mgpuCreateDnVec
- // CHECK: llvm.call @mgpuSpMVBufferSize
- // CHECK: llvm.call @mgpuSpMV
+ // CHECK: %[[BUFFERSIZE_MA:.*]] = llvm.mlir.constant(0 : i32) : i32
+ // CHECK: llvm.call @mgpuSpMVBufferSize(%[[BUFFERSIZE_MA]]
+ // CHECK: %[[SPMV_MA:.*]] = llvm.mlir.constant(0 : i32) : i32
+ // CHECK: llvm.call @mgpuSpMV(%[[SPMV_MA]]
// CHECK: llvm.call @mgpuDestroySpMat
// CHECK: llvm.call @mgpuDestroyDnVec
// CHECK: llvm.call @mgpuStreamSynchronize
@@ -28,6 +30,34 @@ module attributes {gpu.container_module} {
return
}
+ // CHECK-LABEL: func @matvec_transpose
+ // CHECK: llvm.call @mgpuStreamCreate
+ // CHECK: llvm.call @mgpuMemAlloc
+ // CHECK: llvm.call @mgpuMemAlloc
+ // CHECK: llvm.call @mgpuCreateCoo
+ // CHECK: llvm.call @mgpuCreateDnVec
+ // CHECK: %[[BUFFERSIZE_MA:.*]] = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: llvm.call @mgpuSpMVBufferSize(%[[BUFFERSIZE_MA]]
+ // CHECK: %[[SPMV_MA:.*]] = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: llvm.call @mgpuSpMV(%[[SPMV_MA]]
+ // CHECK: llvm.call @mgpuDestroySpMat
+ // CHECK: llvm.call @mgpuDestroyDnVec
+ // CHECK: llvm.call @mgpuStreamSynchronize
+ // CHECK: llvm.call @mgpuStreamDestroy
+ func.func @matvec_transpose(%arg0: index) {
+ %token0 = gpu.wait async
+ %mem1, %token1 = gpu.alloc async [%token0] (%arg0) : memref<?xindex>
+ %mem2, %token2 = gpu.alloc async [%token1] (%arg0) : memref<?xf64>
+ %spmat, %token3 = gpu.create_coo async [%token2] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
+ %dnvec, %token4 = gpu.create_dn_tensor async [%token3] %mem2, %arg0 : index into memref<?xf64>
+ %bufferSz, %token5 = gpu.spmv_buffer_size async [%token4] %spmat{TRANSPOSE}, %dnvec, %dnvec into f64
+ %token6 = gpu.spmv async [%token5] %spmat{TRANSPOSE}, %dnvec, %dnvec, %mem2 : memref<?xf64> into f64
+ %token7 = gpu.destroy_sp_mat async [%token6] %spmat
+ %token8 = gpu.destroy_dn_tensor async [%token7] %dnvec
+ gpu.wait [%token8]
+ return
+ }
+
// CHECK-LABEL: func @matmul
// CHECK: llvm.call @mgpuStreamCreate
// CHECK: llvm.call @mgpuMemAlloc
diff --git a/mlir/test/Dialect/GPU/sparse-roundtrip.mlir b/mlir/test/Dialect/GPU/sparse-roundtrip.mlir
index b88492433ba02e..c3a1168a1f9295 100644
--- a/mlir/test/Dialect/GPU/sparse-roundtrip.mlir
+++ b/mlir/test/Dialect/GPU/sparse-roundtrip.mlir
@@ -28,6 +28,32 @@ module attributes {gpu.container_module} {
return
}
+ // CHECK-LABEL: func @matvec_transpose
+ // CHECK: %{{.*}} = gpu.wait async
+ // CHECK: %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xindex>
+ // CHECK: %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xf64>
+ // CHECK: %{{.*}}, %{{.*}} = gpu.create_coo async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xindex>, memref<?xindex>, memref<?xf64>
+ // CHECK: %{{.*}}, %{{.*}} = gpu.create_dn_tensor async [%{{.*}}] %{{.*}}, %{{.*}} : index into memref<?xf64>
+ // CHECK: %{{.*}}, %{{.*}} = gpu.spmv_buffer_size async [%{{.*}}] %{{.*}}{TRANSPOSE}, %{{.*}}, %{{.*}} into f64
+ // CHECK: %{{.*}} = gpu.spmv async [%{{.*}}] %{{.*}}{TRANSPOSE}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xf64> into f64
+ // CHECK: %{{.*}} = gpu.destroy_sp_mat async [%{{.*}}] %{{.*}}
+ // CHECK: %{{.*}} = gpu.destroy_dn_tensor async [%{{.*}}] %{{.*}}
+ // CHECK: gpu.wait [%{{.*}}]
+ // CHECK: return
+ func.func @matvec_transpose(%arg0: index) {
+ %token0 = gpu.wait async
+ %mem1, %token1 = gpu.alloc async [%token0] (%arg0) : memref<?xindex>
+ %mem2, %token2 = gpu.alloc async [%token1] (%arg0) : memref<?xf64>
+ %spmat, %token3 = gpu.create_coo async [%token2] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
+ %dnvec, %token4 = gpu.create_dn_tensor async [%token3] %mem2, %arg0 : index into memref<?xf64>
+ %bufferSz, %token5 = gpu.spmv_buffer_size async [%token4] %spmat{TRANSPOSE}, %dnvec, %dnvec into f64
+ %token6 = gpu.spmv async [%token5] %spmat{TRANSPOSE}, %dnvec, %dnvec, %mem2 : memref<?xf64> into f64
+ %token7 = gpu.destroy_sp_mat async [%token6] %spmat
+ %token8 = gpu.destroy_dn_tensor async [%token7] %dnvec
+ gpu.wait [%token8]
+ return
+ }
+
// CHECK-LABEL: func @matmul
// CHECK: %{{.*}} = gpu.wait async
// CHECK: %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xindex>
>From 395f14e9b4235f5deb40016a43513de898536922 Mon Sep 17 00:00:00 2001
From: Vito Secona <secona00 at gmail.com>
Date: Fri, 9 Oct 2026 14:28:47 +0700
Subject: [PATCH 2/4] [mlir][sparse] Add libgen support for transposed SpMV
---
.../Transforms/SparseGPUCodegen.cpp | 31 +--
.../GPU/gpu_matvec_transpose_lib.mlir | 53 ++++++
.../GPU/CUDA/sparse-matvec-transpose-lib.mlir | 178 ++++++++++++++++++
3 files changed, 251 insertions(+), 11 deletions(-)
create mode 100644 mlir/test/Dialect/SparseTensor/GPU/gpu_matvec_transpose_lib.mlir
create mode 100644 mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-matvec-transpose-lib.mlir
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
index 8243bf9317c22a..227f8ea42572a1 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
@@ -571,11 +571,12 @@ static Operation *genSpMat(OpBuilder &builder, Location loc,
/// Match and rewrite SpMV kernel.
static LogicalResult rewriteSpMV(PatternRewriter &rewriter,
- linalg::GenericOp op, bool enableRT) {
+ linalg::GenericOp op, bool enableRT,
+ gpu::TransposeMode modeA) {
Location loc = op.getLoc();
Value a = op.getOperand(0);
Value x = op.getOperand(1);
- Value y = op.getOperand(2); // we have y = Ax
+ Value y = op.getOperand(2); // we have y = op(A)x
SmallVector<Value> tokens;
// Only admissible sparse matrix format and dense vectors (no BSR).
@@ -591,8 +592,11 @@ static LogicalResult rewriteSpMV(PatternRewriter &rewriter,
// x : memX -> vecX
// y : memY -> vecY
Value nseA = NumberOfEntriesOp::create(rewriter, loc, a);
- Value szY = linalg::createOrFoldDimOp(rewriter, loc, a, 0);
- Value szX = linalg::createOrFoldDimOp(rewriter, loc, a, 1);
+ Value rows = linalg::createOrFoldDimOp(rewriter, loc, a, 0);
+ Value cols = linalg::createOrFoldDimOp(rewriter, loc, a, 1);
+ bool isTranspose = modeA != gpu::TransposeMode::NON_TRANSPOSE;
+ Value szX = isTranspose ? rows : cols;
+ Value szY = isTranspose ? cols : rows;
Value memR = genFirstPosOrCrds(rewriter, loc, a, format, enableRT);
Value memC = genSecondCrds(rewriter, loc, a, format, enableRT); // or empty
Value memV = ToValuesOp::create(rewriter, loc, a);
@@ -613,7 +617,7 @@ static LogicalResult rewriteSpMV(PatternRewriter &rewriter,
Type tokenTp = rewriter.getType<gpu::AsyncTokenType>();
Value token = genFirstWait(rewriter, loc);
Operation *spGenA =
- genSpMat(rewriter, loc, aTp, spmatHandleTp, tokenTp, token, szY, szX,
+ genSpMat(rewriter, loc, aTp, spmatHandleTp, tokenTp, token, rows, cols,
nseA, rowA, colA, valA, format, enableRT);
Value spMatA = spGenA->getResult(0);
token = spGenA->getResult(1);
@@ -629,7 +633,7 @@ static LogicalResult rewriteSpMV(PatternRewriter &rewriter,
// Precompute buffersize for SpMV.
auto bufferComp = gpu::SpMVBufferSizeOp::create(
- rewriter, loc, indexTp, tokenTp, token, spMatA, dnX, dnY,
+ rewriter, loc, indexTp, tokenTp, token, modeA, spMatA, dnX, dnY,
/*computeType=*/dnYType);
Value bufferSz = bufferComp.getResult(0);
token = bufferComp.getAsyncToken();
@@ -638,9 +642,9 @@ static LogicalResult rewriteSpMV(PatternRewriter &rewriter,
token = buf.getAsyncToken();
// Perform the SpMV.
- auto spmvComp =
- gpu::SpMVOp::create(rewriter, loc, tokenTp, token, spMatA, dnX, dnY,
- /*computeType=*/dnYType, buffer);
+ auto spmvComp = gpu::SpMVOp::create(rewriter, loc, tokenTp, token, modeA,
+ spMatA, dnX, dnY,
+ /*computeType=*/dnYType, buffer);
token = spmvComp.getAsyncToken();
// Copy data back to host and free all the resoures.
@@ -1324,8 +1328,13 @@ struct LinalgOpRewriter : public OpRewritePattern<linalg::GenericOp> {
if (numLoops == 2 && numTensors == 3 &&
linalg::isParallelIterator(iteratorTypes[0]) &&
linalg::isReductionIterator(iteratorTypes[1]) &&
- maps == infer({{i, j}, {j}, {i}}) && matchSumOfMultOfArgs(op)) {
- return rewriteSpMV(rewriter, op, enableRT);
+ matchSumOfMultOfArgs(op)) {
+ if (maps == infer({{i, j}, {j}, {i}}))
+ return rewriteSpMV(rewriter, op, enableRT,
+ gpu::TransposeMode::NON_TRANSPOSE);
+ if (maps == infer({{j, i}, {j}, {i}}))
+ return rewriteSpMV(rewriter, op, enableRT,
+ gpu::TransposeMode::TRANSPOSE);
}
// Recognize a SpGEMM, 2:4-SpMM, or SpMM kernel.
diff --git a/mlir/test/Dialect/SparseTensor/GPU/gpu_matvec_transpose_lib.mlir b/mlir/test/Dialect/SparseTensor/GPU/gpu_matvec_transpose_lib.mlir
new file mode 100644
index 00000000000000..c92a5eb61a4c73
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/GPU/gpu_matvec_transpose_lib.mlir
@@ -0,0 +1,53 @@
+// RUN: mlir-opt %s --sparse-gpu-codegen="num-threads=0" | FileCheck %s
+
+#SortedCOO = #sparse_tensor.encoding<{
+ map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton)
+}>
+
+module {
+
+// Check y += transpose(A) * x. The sparse descriptor keeps A's physical
+// dimensions, while the input and output vector lengths are rows(A) and cols(A).
+
+// CHECK-LABEL: func.func @matvec_transpose(
+// CHECK-SAME: %[[A:.*]]: tensor<?x?xf64, #sparse{{[0-9]*}}>,
+// CHECK-SAME: %[[X:.*]]: tensor<?xf64>,
+// CHECK-SAME: %[[Y:.*]]: tensor<?xf64>) -> tensor<?xf64> {
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[NNZ:.*]] = sparse_tensor.number_of_entries %[[A]]
+// CHECK-DAG: %[[ROWS:.*]] = tensor.dim %[[A]], %[[C0]]
+// CHECK-DAG: %[[COLS:.*]] = tensor.dim %[[A]], %[[C1]]
+// CHECK-DAG: %[[X_HOST:.*]] = bufferization.to_buffer %[[X]]
+// CHECK-DAG: gpu.memcpy async {{\[.*\]}} %[[X_DEVICE:.*]], %[[X_HOST]] :
+// CHECK-DAG: %[[Y_HOST:.*]] = bufferization.to_buffer %[[Y]]
+// CHECK-DAG: gpu.memcpy async {{\[.*\]}} %[[Y_DEVICE:.*]], %[[Y_HOST]] :
+// CHECK: %[[SPMAT:.*]], {{.*}} = gpu.create_coo async {{\[.*\]}} %[[ROWS]], %[[COLS]], %[[NNZ]],
+// CHECK-DAG: %[[DNX:.*]], {{.*}} = gpu.create_dn_tensor async {{\[.*\]}} %[[X_DEVICE]], %[[ROWS]] : index into memref<?xf64>
+// CHECK-DAG: %[[DNY:.*]], {{.*}} = gpu.create_dn_tensor async {{\[.*\]}} %[[Y_DEVICE]], %[[COLS]] : index into memref<?xf64>
+// CHECK: gpu.spmv_buffer_size async {{\[.*\]}} %[[SPMAT]]{TRANSPOSE}, %[[DNX]], %[[DNY]] into f64
+// CHECK: gpu.spmv async {{\[.*\]}} %[[SPMAT]]{TRANSPOSE}, %[[DNX]], %[[DNY]],
+// CHECK: return
+// CHECK: }
+func.func @matvec_transpose(%A: tensor<?x?xf64, #SortedCOO>,
+ %x: tensor<?xf64>,
+ %y_in: tensor<?xf64>) -> tensor<?xf64> {
+ %result = linalg.generic {
+ indexing_maps = [
+ affine_map<(i, j) -> (j, i)>,
+ affine_map<(i, j) -> (j)>,
+ affine_map<(i, j) -> (i)>
+ ],
+ iterator_types = ["parallel", "reduction"]
+ }
+ ins(%A, %x : tensor<?x?xf64, #SortedCOO>, tensor<?xf64>)
+ outs(%y_in : tensor<?xf64>) {
+ ^bb0(%a: f64, %xval: f64, %yval: f64):
+ %product = arith.mulf %a, %xval : f64
+ %sum = arith.addf %yval, %product : f64
+ linalg.yield %sum : f64
+ } -> tensor<?xf64>
+ return %result : tensor<?xf64>
+}
+
+}
diff --git a/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-matvec-transpose-lib.mlir b/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-matvec-transpose-lib.mlir
new file mode 100644
index 00000000000000..c522f17f26bb03
--- /dev/null
+++ b/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-matvec-transpose-lib.mlir
@@ -0,0 +1,178 @@
+// NOTE: this test requires gpu-sm80
+//
+// DEFINE: %{compile} = mlir-opt %s \
+// DEFINE: --sparsifier="enable-gpu-libgen gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71 gpu-format=%gpu_compilation_format
+// DEFINE: %{run} = mlir-runner \
+// DEFINE: --shared-libs=%mlir_cuda_runtime \
+// DEFINE: --shared-libs=%mlir_c_runner_utils \
+// DEFINE: --e main --entry-point-result=void \
+// DEFINE: | FileCheck %s
+//
+// with RT lib (SoA COO):
+//
+// RUN: %{compile} enable-runtime-library=true" | %{run}
+//
+// without RT lib (AoS COO): note, may fall back to CPU
+//
+// RUN: %{compile} enable-runtime-library=false" | %{run}
+
+#SortedCOO = #sparse_tensor.encoding<{
+ map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton)
+}>
+
+#CSR = #sparse_tensor.encoding<{
+ map = (d0, d1) -> (d0 : dense, d1 : compressed),
+ posWidth = 32,
+ crdWidth = 32
+}>
+
+#CSC = #sparse_tensor.encoding<{
+ map = (d0, d1) -> (d1 : dense, d0 : compressed),
+ posWidth = 64,
+ crdWidth = 64
+}>
+
+#matvec_transpose = {
+ indexing_maps = [
+ affine_map<(i, j) -> (j, i)>,
+ affine_map<(i, j) -> (j)>,
+ affine_map<(i, j) -> (i)>
+ ],
+ iterator_types = ["parallel", "reduction"]
+}
+
+module {
+ llvm.func @mgpuCreateSparseEnv()
+ llvm.func @mgpuDestroySparseEnv()
+
+ // Compute matrix vector y = A^T x on COO with default index coordinates.
+ func.func @matvecCOO(%A: tensor<?x?xf64, #SortedCOO>, %x: tensor<?xf64>, %y_in: tensor<?xf64>) -> tensor<?xf64> {
+ %y_out = linalg.generic #matvec_transpose
+ ins(%A, %x: tensor<?x?xf64, #SortedCOO>, tensor<?xf64>)
+ outs(%y_in: tensor<?xf64>) {
+ ^bb0(%a: f64, %xval: f64, %yval: f64):
+ %product = arith.mulf %a, %xval : f64
+ %sum = arith.addf %yval, %product : f64
+ linalg.yield %sum : f64
+ } -> tensor<?xf64>
+ return %y_out : tensor<?xf64>
+ }
+
+ // Compute matrix vector y = A^T x on CSR with 32-bit positions and coordinates.
+ func.func @matvecCSR(%A: tensor<?x?xf64, #CSR>, %x: tensor<?xf64>, %y_in: tensor<?xf64>) -> tensor<?xf64> {
+ %y_out = linalg.generic #matvec_transpose
+ ins(%A, %x: tensor<?x?xf64, #CSR>, tensor<?xf64>)
+ outs(%y_in: tensor<?xf64>) {
+ ^bb0(%a: f64, %xval: f64, %yval: f64):
+ %product = arith.mulf %a, %xval : f64
+ %sum = arith.addf %yval, %product : f64
+ linalg.yield %sum : f64
+ } -> tensor<?xf64>
+ return %y_out : tensor<?xf64>
+ }
+
+ // Compute matrix vector y = A^T x on CSC with 64-bit positions and coordinates.
+ func.func @matvecCSC(%A: tensor<?x?xf64, #CSC>, %x: tensor<?xf64>, %y_in: tensor<?xf64>) -> tensor<?xf64> {
+ %y_out = linalg.generic #matvec_transpose
+ ins(%A, %x: tensor<?x?xf64, #CSC>, tensor<?xf64>)
+ outs(%y_in: tensor<?xf64>) {
+ ^bb0(%a: f64, %xval: f64, %yval: f64):
+ %product = arith.mulf %a, %xval : f64
+ %sum = arith.addf %yval, %product : f64
+ linalg.yield %sum : f64
+ } -> tensor<?xf64>
+ return %y_out : tensor<?xf64>
+ }
+
+ func.func @main() {
+ llvm.call @mgpuCreateSparseEnv() : () -> ()
+ %f0 = arith.constant 0.0 : f64
+ %f1 = arith.constant 1.0 : f64
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+
+ // Stress test with a rectangular dense matrix DA.
+ %DA = tensor.generate {
+ ^bb0(%i: index, %j: index):
+ %k = arith.addi %i, %j : index
+ %l = arith.index_cast %k : index to i64
+ %f = arith.uitofp %l : i64 to f64
+ tensor.yield %f : f64
+ } : tensor<32x64xf64>
+
+ // Convert to a "sparse" m x n matrix A.
+ %Acoo = sparse_tensor.convert %DA : tensor<32x64xf64> to tensor<?x?xf64, #SortedCOO>
+ %Acsr = sparse_tensor.convert %DA : tensor<32x64xf64> to tensor<?x?xf64, #CSR>
+ %Acsc = sparse_tensor.convert %DA : tensor<32x64xf64> to tensor<?x?xf64, #CSC>
+
+ // Initialize dense vector with m elements:
+ // (1, 2, 3, 4, ..., m)
+ %d0 = tensor.dim %Acoo, %c0 : tensor<?x?xf64, #SortedCOO>
+ %x = tensor.generate %d0 {
+ ^bb0(%i : index):
+ %k = arith.addi %i, %c1 : index
+ %j = arith.index_cast %k : index to i64
+ %f = arith.uitofp %j : i64 to f64
+ tensor.yield %f : f64
+ } : tensor<?xf64>
+
+ // Initialize dense vectors to n zeros and n ones.
+ %d1 = tensor.dim %Acoo, %c1 : tensor<?x?xf64, #SortedCOO>
+ %y0 = tensor.generate %d1 {
+ ^bb0(%i : index):
+ tensor.yield %f0 : f64
+ } : tensor<?xf64>
+ %y1 = tensor.generate %d1 {
+ ^bb0(%i : index):
+ tensor.yield %f1 : f64
+ } : tensor<?xf64>
+
+ // Call the kernels.
+ %0 = call @matvecCOO(%Acoo, %x, %y0) : (tensor<?x?xf64, #SortedCOO>,
+ tensor<?xf64>,
+ tensor<?xf64>) -> tensor<?xf64>
+ %1 = call @matvecCSR(%Acsr, %x, %y0) : (tensor<?x?xf64, #CSR>,
+ tensor<?xf64>,
+ tensor<?xf64>) -> tensor<?xf64>
+ %2 = call @matvecCSC(%Acsc, %x, %y0) : (tensor<?x?xf64, #CSC>,
+ tensor<?xf64>,
+ tensor<?xf64>) -> tensor<?xf64>
+ %3 = call @matvecCOO(%Acoo, %x, %y1) : (tensor<?x?xf64, #SortedCOO>,
+ tensor<?xf64>,
+ tensor<?xf64>) -> tensor<?xf64>
+ %4 = call @matvecCSR(%Acsr, %x, %y1) : (tensor<?x?xf64, #CSR>,
+ tensor<?xf64>,
+ tensor<?xf64>) -> tensor<?xf64>
+ %5 = call @matvecCSC(%Acsc, %x, %y1) : (tensor<?x?xf64, #CSC>,
+ tensor<?xf64>,
+ tensor<?xf64>) -> tensor<?xf64>
+
+ //
+ // Sanity check on the results.
+ //
+ // CHECK-COUNT-3: ( 10912, 11440, 11968, 12496, 13024, 13552, 14080, 14608, 15136, 15664, 16192, 16720, 17248, 17776, 18304, 18832, 19360, 19888, 20416, 20944, 21472, 22000, 22528, 23056, 23584, 24112, 24640, 25168, 25696, 26224, 26752, 27280, 27808, 28336, 28864, 29392, 29920, 30448, 30976, 31504, 32032, 32560, 33088, 33616, 34144, 34672, 35200, 35728, 36256, 36784, 37312, 37840, 38368, 38896, 39424, 39952, 40480, 41008, 41536, 42064, 42592, 43120, 43648, 44176 )
+ //
+ // CHECK-COUNT-3: ( 10913, 11441, 11969, 12497, 13025, 13553, 14081, 14609, 15137, 15665, 16193, 16721, 17249, 17777, 18305, 18833, 19361, 19889, 20417, 20945, 21473, 22001, 22529, 23057, 23585, 24113, 24641, 25169, 25697, 26225, 26753, 27281, 27809, 28337, 28865, 29393, 29921, 30449, 30977, 31505, 32033, 32561, 33089, 33617, 34145, 34673, 35201, 35729, 36257, 36785, 37313, 37841, 38369, 38897, 39425, 39953, 40481, 41009, 41537, 42065, 42593, 43121, 43649, 44177 )
+ //
+ %pb0 = vector.transfer_read %0[%c0], %f0 : tensor<?xf64>, vector<64xf64>
+ vector.print %pb0 : vector<64xf64>
+ %pb1 = vector.transfer_read %1[%c0], %f0 : tensor<?xf64>, vector<64xf64>
+ vector.print %pb1 : vector<64xf64>
+ %pb2 = vector.transfer_read %2[%c0], %f0 : tensor<?xf64>, vector<64xf64>
+ vector.print %pb2 : vector<64xf64>
+ %pb3 = vector.transfer_read %3[%c0], %f0 : tensor<?xf64>, vector<64xf64>
+ vector.print %pb3 : vector<64xf64>
+ %pb4 = vector.transfer_read %4[%c0], %f0 : tensor<?xf64>, vector<64xf64>
+ vector.print %pb4 : vector<64xf64>
+ %pb5 = vector.transfer_read %5[%c0], %f0 : tensor<?xf64>, vector<64xf64>
+ vector.print %pb5 : vector<64xf64>
+
+ // Release the resources.
+ bufferization.dealloc_tensor %Acoo : tensor<?x?xf64, #SortedCOO>
+ bufferization.dealloc_tensor %Acsr : tensor<?x?xf64, #CSR>
+ bufferization.dealloc_tensor %Acsc : tensor<?x?xf64, #CSC>
+
+ llvm.call @mgpuDestroySparseEnv() : () -> ()
+ return
+ }
+}
>From 8359a2f5c1b051417097f622b5943bd7423bfa97 Mon Sep 17 00:00:00 2001
From: Vito Secona <secona00 at gmail.com>
Date: Fri, 9 Oct 2026 14:57:00 +0700
Subject: [PATCH 3/4] [mlir][gpu] Test transposed SpMM roundtrip and runtime
lowering
---
.../lower-sparse-to-gpu-runtime-calls.mlir | 75 ++++++++++++++++++
mlir/test/Dialect/GPU/sparse-roundtrip.mlir | 78 +++++++++++++++++++
2 files changed, 153 insertions(+)
diff --git a/mlir/test/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir b/mlir/test/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir
index 2e45813dc1e1c6..9dcdabe557065c 100644
--- a/mlir/test/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir
+++ b/mlir/test/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir
@@ -84,6 +84,81 @@ module attributes {gpu.container_module} {
return
}
+ // CHECK-LABEL: func @matmul_transpose_a
+ // CHECK: llvm.call @mgpuCreateCoo
+ // CHECK: llvm.call @mgpuCreateDnMat
+ // CHECK: %[[BUFFERSIZE_MA:.*]] = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: %[[BUFFERSIZE_MB:.*]] = llvm.mlir.constant(0 : i32) : i32
+ // CHECK: llvm.call @mgpuSpMMBufferSize(%[[BUFFERSIZE_MA]], %[[BUFFERSIZE_MB]],
+ // CHECK: %[[SPMM_MA:.*]] = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: %[[SPMM_MB:.*]] = llvm.mlir.constant(0 : i32) : i32
+ // CHECK: llvm.call @mgpuSpMM(%[[SPMM_MA]], %[[SPMM_MB]],
+ // CHECK: llvm.call @mgpuDestroySpMat
+ // CHECK: llvm.call @mgpuDestroyDnMat
+ func.func @matmul_transpose_a(%arg0: index) {
+ %token0 = gpu.wait async
+ %mem1, %token1 = gpu.alloc async [%token0] (%arg0) : memref<?xindex>
+ %mem2, %token2 = gpu.alloc async [%token1] (%arg0) : memref<?xf64>
+ %spmat, %token3 = gpu.create_coo async [%token2] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
+ %dnmat, %token4 = gpu.create_dn_tensor async [%token3] %mem2, %arg0, %arg0 : index, index into memref<?xf64>
+ %bufferSz, %token5 = gpu.spmm_buffer_size async [%token4] %spmat{TRANSPOSE}, %dnmat, %dnmat : index into f64
+ %token6 = gpu.spmm async [%token5] %spmat{TRANSPOSE}, %dnmat, %dnmat, %mem2 : memref<?xf64> into f64
+ %token7 = gpu.destroy_sp_mat async [%token6] %spmat
+ %token8 = gpu.destroy_dn_tensor async [%token7] %dnmat
+ gpu.wait [%token8]
+ return
+ }
+
+ // CHECK-LABEL: func @matmul_transpose_b
+ // CHECK: llvm.call @mgpuCreateCoo
+ // CHECK: llvm.call @mgpuCreateDnMat
+ // CHECK: %[[BUFFERSIZE_MA:.*]] = llvm.mlir.constant(0 : i32) : i32
+ // CHECK: %[[BUFFERSIZE_MB:.*]] = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: llvm.call @mgpuSpMMBufferSize(%[[BUFFERSIZE_MA]], %[[BUFFERSIZE_MB]],
+ // CHECK: %[[SPMM_MA:.*]] = llvm.mlir.constant(0 : i32) : i32
+ // CHECK: %[[SPMM_MB:.*]] = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: llvm.call @mgpuSpMM(%[[SPMM_MA]], %[[SPMM_MB]],
+ // CHECK: llvm.call @mgpuDestroySpMat
+ // CHECK: llvm.call @mgpuDestroyDnMat
+ func.func @matmul_transpose_b(%arg0: index) {
+ %token0 = gpu.wait async
+ %mem1, %token1 = gpu.alloc async [%token0] (%arg0) : memref<?xindex>
+ %mem2, %token2 = gpu.alloc async [%token1] (%arg0) : memref<?xf64>
+ %spmat, %token3 = gpu.create_coo async [%token2] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
+ %dnmat, %token4 = gpu.create_dn_tensor async [%token3] %mem2, %arg0, %arg0 : index, index into memref<?xf64>
+ %bufferSz, %token5 = gpu.spmm_buffer_size async [%token4] %spmat, %dnmat{TRANSPOSE}, %dnmat : index into f64
+ %token6 = gpu.spmm async [%token5] %spmat, %dnmat{TRANSPOSE}, %dnmat, %mem2 : memref<?xf64> into f64
+ %token7 = gpu.destroy_sp_mat async [%token6] %spmat
+ %token8 = gpu.destroy_dn_tensor async [%token7] %dnmat
+ gpu.wait [%token8]
+ return
+ }
+
+ // CHECK-LABEL: func @matmul_transpose_ab
+ // CHECK: llvm.call @mgpuCreateCoo
+ // CHECK: llvm.call @mgpuCreateDnMat
+ // CHECK: %[[BUFFERSIZE_MA:.*]] = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: %[[BUFFERSIZE_MB:.*]] = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: llvm.call @mgpuSpMMBufferSize(%[[BUFFERSIZE_MA]], %[[BUFFERSIZE_MB]],
+ // CHECK: %[[SPMM_MA:.*]] = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: %[[SPMM_MB:.*]] = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: llvm.call @mgpuSpMM(%[[SPMM_MA]], %[[SPMM_MB]],
+ // CHECK: llvm.call @mgpuDestroySpMat
+ // CHECK: llvm.call @mgpuDestroyDnMat
+ func.func @matmul_transpose_ab(%arg0: index) {
+ %token0 = gpu.wait async
+ %mem1, %token1 = gpu.alloc async [%token0] (%arg0) : memref<?xindex>
+ %mem2, %token2 = gpu.alloc async [%token1] (%arg0) : memref<?xf64>
+ %spmat, %token3 = gpu.create_coo async [%token2] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
+ %dnmat, %token4 = gpu.create_dn_tensor async [%token3] %mem2, %arg0, %arg0 : index, index into memref<?xf64>
+ %bufferSz, %token5 = gpu.spmm_buffer_size async [%token4] %spmat{TRANSPOSE}, %dnmat{TRANSPOSE}, %dnmat : index into f64
+ %token6 = gpu.spmm async [%token5] %spmat{TRANSPOSE}, %dnmat{TRANSPOSE}, %dnmat, %mem2 : memref<?xf64> into f64
+ %token7 = gpu.destroy_sp_mat async [%token6] %spmat
+ %token8 = gpu.destroy_dn_tensor async [%token7] %dnmat
+ gpu.wait [%token8]
+ return
+ }
+
// CHECK-LABEL: func @spgemm
// CHECK: llvm.call @mgpuStreamCreate
// CHECK: llvm.call @mgpuMemAlloc
diff --git a/mlir/test/Dialect/GPU/sparse-roundtrip.mlir b/mlir/test/Dialect/GPU/sparse-roundtrip.mlir
index c3a1168a1f9295..81433f9bcbcc6b 100644
--- a/mlir/test/Dialect/GPU/sparse-roundtrip.mlir
+++ b/mlir/test/Dialect/GPU/sparse-roundtrip.mlir
@@ -80,6 +80,84 @@ module attributes {gpu.container_module} {
return
}
+ // CHECK-LABEL: func @matmul_transpose_a
+ // CHECK: %{{.*}} = gpu.wait async
+ // CHECK: %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xindex>
+ // CHECK: %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xf64>
+ // CHECK: %{{.*}}, %{{.*}} = gpu.create_coo async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xindex>, memref<?xindex>, memref<?xf64>
+ // CHECK: %{{.*}}, %{{.*}} = gpu.create_dn_tensor async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}} : index, index into memref<?xf64>
+ // CHECK: %{{.*}}, %{{.*}} = gpu.spmm_buffer_size async [%{{.*}}] %{{.*}}{TRANSPOSE}, %{{.*}}, %{{.*}} : index into f64
+ // CHECK: %{{.*}} = gpu.spmm async [%{{.*}}] %{{.*}}{TRANSPOSE}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xf64> into f64
+ // CHECK: %{{.*}} = gpu.destroy_sp_mat async [%{{.*}}] %{{.*}}
+ // CHECK: %{{.*}} = gpu.destroy_dn_tensor async [%{{.*}}] %{{.*}}
+ // CHECK: gpu.wait [%{{.*}}]
+ // CHECK: return
+ func.func @matmul_transpose_a(%arg0: index) {
+ %token0 = gpu.wait async
+ %mem1, %token1 = gpu.alloc async [%token0] (%arg0) : memref<?xindex>
+ %mem2, %token2 = gpu.alloc async [%token1] (%arg0) : memref<?xf64>
+ %spmat, %token3 = gpu.create_coo async [%token2] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
+ %dnmat, %token4 = gpu.create_dn_tensor async [%token3] %mem2, %arg0, %arg0 : index, index into memref<?xf64>
+ %bufferSz, %token5 = gpu.spmm_buffer_size async [%token4] %spmat{TRANSPOSE}, %dnmat, %dnmat : index into f64
+ %token6 = gpu.spmm async [%token5] %spmat{TRANSPOSE}, %dnmat, %dnmat, %mem2 : memref<?xf64> into f64
+ %token7 = gpu.destroy_sp_mat async [%token6] %spmat
+ %token8 = gpu.destroy_dn_tensor async [%token7] %dnmat
+ gpu.wait [%token8]
+ return
+ }
+
+ // CHECK-LABEL: func @matmul_transpose_b
+ // CHECK: %{{.*}} = gpu.wait async
+ // CHECK: %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xindex>
+ // CHECK: %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xf64>
+ // CHECK: %{{.*}}, %{{.*}} = gpu.create_coo async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xindex>, memref<?xindex>, memref<?xf64>
+ // CHECK: %{{.*}}, %{{.*}} = gpu.create_dn_tensor async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}} : index, index into memref<?xf64>
+ // CHECK: %{{.*}}, %{{.*}} = gpu.spmm_buffer_size async [%{{.*}}] %{{.*}}, %{{.*}}{TRANSPOSE}, %{{.*}} : index into f64
+ // CHECK: %{{.*}} = gpu.spmm async [%{{.*}}] %{{.*}}, %{{.*}}{TRANSPOSE}, %{{.*}}, %{{.*}} : memref<?xf64> into f64
+ // CHECK: %{{.*}} = gpu.destroy_sp_mat async [%{{.*}}] %{{.*}}
+ // CHECK: %{{.*}} = gpu.destroy_dn_tensor async [%{{.*}}] %{{.*}}
+ // CHECK: gpu.wait [%{{.*}}]
+ // CHECK: return
+ func.func @matmul_transpose_b(%arg0: index) {
+ %token0 = gpu.wait async
+ %mem1, %token1 = gpu.alloc async [%token0] (%arg0) : memref<?xindex>
+ %mem2, %token2 = gpu.alloc async [%token1] (%arg0) : memref<?xf64>
+ %spmat, %token3 = gpu.create_coo async [%token2] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
+ %dnmat, %token4 = gpu.create_dn_tensor async [%token3] %mem2, %arg0, %arg0 : index, index into memref<?xf64>
+ %bufferSz, %token5 = gpu.spmm_buffer_size async [%token4] %spmat, %dnmat{TRANSPOSE}, %dnmat : index into f64
+ %token6 = gpu.spmm async [%token5] %spmat, %dnmat{TRANSPOSE}, %dnmat, %mem2 : memref<?xf64> into f64
+ %token7 = gpu.destroy_sp_mat async [%token6] %spmat
+ %token8 = gpu.destroy_dn_tensor async [%token7] %dnmat
+ gpu.wait [%token8]
+ return
+ }
+
+ // CHECK-LABEL: func @matmul_transpose_ab
+ // CHECK: %{{.*}} = gpu.wait async
+ // CHECK: %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xindex>
+ // CHECK: %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xf64>
+ // CHECK: %{{.*}}, %{{.*}} = gpu.create_coo async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<?xindex>, memref<?xindex>, memref<?xf64>
+ // CHECK: %{{.*}}, %{{.*}} = gpu.create_dn_tensor async [%{{.*}}] %{{.*}}, %{{.*}}, %{{.*}} : index, index into memref<?xf64>
+ // CHECK: %{{.*}}, %{{.*}} = gpu.spmm_buffer_size async [%{{.*}}] %{{.*}}{TRANSPOSE}, %{{.*}}{TRANSPOSE}, %{{.*}} : index into f64
+ // CHECK: %{{.*}} = gpu.spmm async [%{{.*}}] %{{.*}}{TRANSPOSE}, %{{.*}}{TRANSPOSE}, %{{.*}}, %{{.*}} : memref<?xf64> into f64
+ // CHECK: %{{.*}} = gpu.destroy_sp_mat async [%{{.*}}] %{{.*}}
+ // CHECK: %{{.*}} = gpu.destroy_dn_tensor async [%{{.*}}] %{{.*}}
+ // CHECK: gpu.wait [%{{.*}}]
+ // CHECK: return
+ func.func @matmul_transpose_ab(%arg0: index) {
+ %token0 = gpu.wait async
+ %mem1, %token1 = gpu.alloc async [%token0] (%arg0) : memref<?xindex>
+ %mem2, %token2 = gpu.alloc async [%token1] (%arg0) : memref<?xf64>
+ %spmat, %token3 = gpu.create_coo async [%token2] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>
+ %dnmat, %token4 = gpu.create_dn_tensor async [%token3] %mem2, %arg0, %arg0 : index, index into memref<?xf64>
+ %bufferSz, %token5 = gpu.spmm_buffer_size async [%token4] %spmat{TRANSPOSE}, %dnmat{TRANSPOSE}, %dnmat : index into f64
+ %token6 = gpu.spmm async [%token5] %spmat{TRANSPOSE}, %dnmat{TRANSPOSE}, %dnmat, %mem2 : memref<?xf64> into f64
+ %token7 = gpu.destroy_sp_mat async [%token6] %spmat
+ %token8 = gpu.destroy_dn_tensor async [%token7] %dnmat
+ gpu.wait [%token8]
+ return
+ }
+
// CHECK-LABEL: func @spgemm
// CHECK: %{{.*}} = gpu.wait async
// CHECK: %{{.*}}, %{{.*}} = gpu.alloc async [%{{.*}}] (%{{.*}}) : memref<?xindex>
>From fa7f226771d3cb3d17c0092cd42af516c1c4451e Mon Sep 17 00:00:00 2001
From: Vito Secona <secona00 at gmail.com>
Date: Fri, 9 Oct 2026 15:14:33 +0700
Subject: [PATCH 4/4] [mlir][sparse] Add libgen support for transposed SpMM
---
.../Transforms/SparseGPUCodegen.cpp | 62 ++-
.../GPU/gpu_matmul_transpose_lib.mlir | 157 ++++++
.../GPU/CUDA/sparse-matmul-transpose-lib.mlir | 456 ++++++++++++++++++
3 files changed, 656 insertions(+), 19 deletions(-)
create mode 100644 mlir/test/Dialect/SparseTensor/GPU/gpu_matmul_transpose_lib.mlir
create mode 100644 mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-matmul-transpose-lib.mlir
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
index 227f8ea42572a1..684876597eb95d 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
@@ -673,11 +673,13 @@ static LogicalResult rewriteSpMV(PatternRewriter &rewriter,
/// Match and rewrite SpMM kernel.
static LogicalResult rewriteSpMM(PatternRewriter &rewriter,
- linalg::GenericOp op, bool enableRT) {
+ linalg::GenericOp op, bool enableRT,
+ gpu::TransposeMode modeA,
+ gpu::TransposeMode modeB) {
Location loc = op.getLoc();
Value a = op.getOperand(0);
Value b = op.getOperand(1);
- Value c = op.getOperand(2); // we have C = AB
+ Value c = op.getOperand(2); // we have C = op(A)op(B)
SmallVector<Value> tokens;
// Only admissible sparse matrix format and dense matrices (no BSR).
@@ -693,9 +695,15 @@ static LogicalResult rewriteSpMM(PatternRewriter &rewriter,
// b : bufB -> matB
// c : bufC -> matC
Value nseA = NumberOfEntriesOp::create(rewriter, loc, a);
- 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 rows = linalg::createOrFoldDimOp(rewriter, loc, a, 0);
+ Value cols = linalg::createOrFoldDimOp(rewriter, loc, a, 1);
+ bool isTransposeA = modeA != gpu::TransposeMode::NON_TRANSPOSE;
+ bool isTransposeB = modeB != gpu::TransposeMode::NON_TRANSPOSE;
+ Value szm = isTransposeA ? cols : rows;
+ Value szk = isTransposeA ? rows : cols;
+ Value szn = linalg::createOrFoldDimOp(rewriter, loc, b, isTransposeB ? 0 : 1);
+ SmallVector<Value> bSizes = isTransposeB ? SmallVector<Value>{szn, szk}
+ : SmallVector<Value>{szk, szn};
Value memR = genFirstPosOrCrds(rewriter, loc, a, format, enableRT);
Value memC = genSecondCrds(rewriter, loc, a, format, enableRT); // or empty
Value memV = ToValuesOp::create(rewriter, loc, a);
@@ -716,13 +724,12 @@ static LogicalResult rewriteSpMM(PatternRewriter &rewriter,
Type tokenTp = rewriter.getType<gpu::AsyncTokenType>();
Value token = genFirstWait(rewriter, loc);
Operation *spGenA =
- genSpMat(rewriter, loc, aTp, spMatHandleTp, tokenTp, token, szm, szk,
+ genSpMat(rewriter, loc, aTp, spMatHandleTp, tokenTp, token, rows, cols,
nseA, rowA, colA, valA, format, enableRT);
Value spMatA = spGenA->getResult(0);
token = spGenA->getResult(1);
- auto dmatB =
- gpu::CreateDnTensorOp::create(rewriter, loc, dnTensorHandleTp, tokenTp,
- token, matB, SmallVector<Value>{szk, szn});
+ auto dmatB = gpu::CreateDnTensorOp::create(rewriter, loc, dnTensorHandleTp,
+ tokenTp, token, matB, bSizes);
Value dnB = dmatB.getResult(0);
token = dmatB.getAsyncToken();
auto dmatC =
@@ -734,7 +741,7 @@ static LogicalResult rewriteSpMM(PatternRewriter &rewriter,
// Precompute buffersize for SpMM.
auto bufferComp = gpu::SpMMBufferSizeOp::create(
- rewriter, loc, indexTp, tokenTp, token, spMatA, dnB, dnC,
+ rewriter, loc, indexTp, tokenTp, token, modeA, modeB, spMatA, dnB, dnC,
/*computeType=*/dmatCType);
Value bufferSz = bufferComp.getResult(0);
token = bufferComp.getAsyncToken();
@@ -744,9 +751,9 @@ static LogicalResult rewriteSpMM(PatternRewriter &rewriter,
auto dnCType = llvm::cast<ShapedType>(c.getType()).getElementType();
// Perform the SpMM.
- auto spmmComp =
- gpu::SpMMOp::create(rewriter, loc, tokenTp, token, spMatA, dnB, dnC,
- /*computeType=*/dnCType, buffer);
+ auto spmmComp = gpu::SpMMOp::create(rewriter, loc, tokenTp, token, modeA,
+ modeB, spMatA, dnB, dnC,
+ /*computeType=*/dnCType, buffer);
token = spmmComp.getAsyncToken();
// Copy data back to host and free all the resoures.
@@ -1342,12 +1349,29 @@ struct LinalgOpRewriter : public OpRewritePattern<linalg::GenericOp> {
linalg::isParallelIterator(iteratorTypes[0]) &&
linalg::isParallelIterator(iteratorTypes[1]) &&
linalg::isReductionIterator(iteratorTypes[2]) &&
- maps == infer({{i, k}, {k, j}, {i, j}}) && matchSumOfMultOfArgs(op)) {
- if (!isDenseTensor(op.getOperand(0)) && !isDenseTensor(op.getOperand(1)))
- return rewriteSpGEMM(rewriter, op, enableRT);
- if (isConversionInto24(op.getOperand(0)))
- return rewrite2To4SpMM(rewriter, op);
- return rewriteSpMM(rewriter, op, enableRT);
+ matchSumOfMultOfArgs(op)) {
+ if (maps == infer({{i, k}, {k, j}, {i, j}})) {
+ if (!isDenseTensor(op.getOperand(0)) &&
+ !isDenseTensor(op.getOperand(1)))
+ return rewriteSpGEMM(rewriter, op, enableRT);
+ if (isConversionInto24(op.getOperand(0)))
+ return rewrite2To4SpMM(rewriter, op);
+ return rewriteSpMM(rewriter, op, enableRT,
+ gpu::TransposeMode::NON_TRANSPOSE,
+ gpu::TransposeMode::NON_TRANSPOSE);
+ }
+ if (maps == infer({{k, i}, {k, j}, {i, j}}))
+ return rewriteSpMM(rewriter, op, enableRT,
+ gpu::TransposeMode::TRANSPOSE,
+ gpu::TransposeMode::NON_TRANSPOSE);
+ if (maps == infer({{i, k}, {j, k}, {i, j}}))
+ return rewriteSpMM(rewriter, op, enableRT,
+ gpu::TransposeMode::NON_TRANSPOSE,
+ gpu::TransposeMode::TRANSPOSE);
+ if (maps == infer({{k, i}, {j, k}, {i, j}}))
+ return rewriteSpMM(rewriter, op, enableRT,
+ gpu::TransposeMode::TRANSPOSE,
+ gpu::TransposeMode::TRANSPOSE);
}
// Recognize a SDDMM kernel.
diff --git a/mlir/test/Dialect/SparseTensor/GPU/gpu_matmul_transpose_lib.mlir b/mlir/test/Dialect/SparseTensor/GPU/gpu_matmul_transpose_lib.mlir
new file mode 100644
index 00000000000000..fb333066310a92
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/GPU/gpu_matmul_transpose_lib.mlir
@@ -0,0 +1,157 @@
+// RUN: mlir-opt %s --linalg-generalize-named-ops --sparse-gpu-codegen="num-threads=0" | FileCheck %s --enable-var-scope
+
+#SortedCOO = #sparse_tensor.encoding<{
+ map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton)
+}>
+
+// CHECK-LABEL: func.func @matmul_transpose_a(
+// CHECK-SAME: %[[A:.*]]: tensor<?x?xf64, #sparse{{[0-9]*}}>,
+// CHECK-SAME: %[[B:.*]]: tensor<?x?xf64>,
+// CHECK-SAME: %[[C:.*]]: tensor<?x?xf64>) -> tensor<?x?xf64> {
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[NNZ:.*]] = sparse_tensor.number_of_entries %[[A]]
+// CHECK-DAG: %[[ROWS:.*]] = tensor.dim %[[A]], %[[C0]]
+// CHECK-DAG: %[[COLS:.*]] = tensor.dim %[[A]], %[[C1]]
+// CHECK-DAG: %[[N:.*]] = tensor.dim %[[B]], %[[C1]]
+// CHECK-DAG: %[[B_HOST:.*]] = bufferization.to_buffer %[[B]]
+// CHECK-DAG: gpu.memcpy async {{\[.*\]}} %[[B_DEVICE:.*]], %[[B_HOST]] :
+// CHECK-DAG: %[[C_HOST:.*]] = bufferization.to_buffer %[[C]]
+// CHECK-DAG: gpu.memcpy async {{\[.*\]}} %[[C_DEVICE:.*]], %[[C_HOST]] :
+// CHECK: %[[SPMAT:.*]], {{.*}} = gpu.create_coo async {{\[.*\]}} %[[ROWS]], %[[COLS]], %[[NNZ]],
+// CHECK: %[[DNB:.*]], {{.*}} = gpu.create_dn_tensor async {{\[.*\]}} %[[B_DEVICE]], %[[ROWS]], %[[N]] : index, index into memref<?x?xf64>
+// CHECK: %[[DNC:.*]], {{.*}} = gpu.create_dn_tensor async {{\[.*\]}} %[[C_DEVICE]], %[[COLS]], %[[N]] : index, index into memref<?x?xf64>
+// CHECK: gpu.spmm_buffer_size async {{\[.*\]}} %[[SPMAT]]{TRANSPOSE}, %[[DNB]], %[[DNC]] : index into f64
+// CHECK: gpu.spmm async {{\[.*\]}} %[[SPMAT]]{TRANSPOSE}, %[[DNB]], %[[DNC]],
+// CHECK: return
+// CHECK: }
+func.func @matmul_transpose_a(%A: tensor<?x?xf64, #SortedCOO>,
+ %B: tensor<?x?xf64>,
+ %C_in: tensor<?x?xf64>) -> tensor<?x?xf64> {
+ %C_out = linalg.matmul
+ indexing_maps = [
+ affine_map<(i, j, k) -> (k, i)>,
+ affine_map<(i, j, k) -> (k, j)>,
+ affine_map<(i, j, k) -> (i, j)>
+ ]
+ ins(%A, %B: tensor<?x?xf64, #SortedCOO>, tensor<?x?xf64>)
+ outs(%C_in: tensor<?x?xf64>) -> tensor<?x?xf64>
+ return %C_out : tensor<?x?xf64>
+}
+
+// CHECK-LABEL: func.func @matmul_transpose_b(
+// CHECK-SAME: %[[A:.*]]: tensor<?x?xf64, #sparse{{[0-9]*}}>,
+// CHECK-SAME: %[[B:.*]]: tensor<?x?xf64>,
+// CHECK-SAME: %[[C:.*]]: tensor<?x?xf64>) -> tensor<?x?xf64> {
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[NNZ:.*]] = sparse_tensor.number_of_entries %[[A]]
+// CHECK-DAG: %[[ROWS:.*]] = tensor.dim %[[A]], %[[C0]]
+// CHECK-DAG: %[[COLS:.*]] = tensor.dim %[[A]], %[[C1]]
+// CHECK-DAG: %[[N:.*]] = tensor.dim %[[B]], %[[C0]]
+// CHECK-DAG: %[[B_HOST:.*]] = bufferization.to_buffer %[[B]]
+// CHECK-DAG: gpu.memcpy async {{\[.*\]}} %[[B_DEVICE:.*]], %[[B_HOST]] :
+// CHECK-DAG: %[[C_HOST:.*]] = bufferization.to_buffer %[[C]]
+// CHECK-DAG: gpu.memcpy async {{\[.*\]}} %[[C_DEVICE:.*]], %[[C_HOST]] :
+// CHECK: %[[SPMAT:.*]], {{.*}} = gpu.create_coo async {{\[.*\]}} %[[ROWS]], %[[COLS]], %[[NNZ]],
+// CHECK: %[[DNB:.*]], {{.*}} = gpu.create_dn_tensor async {{\[.*\]}} %[[B_DEVICE]], %[[N]], %[[COLS]] : index, index into memref<?x?xf64>
+// CHECK: %[[DNC:.*]], {{.*}} = gpu.create_dn_tensor async {{\[.*\]}} %[[C_DEVICE]], %[[ROWS]], %[[N]] : index, index into memref<?x?xf64>
+// CHECK: gpu.spmm_buffer_size async {{\[.*\]}} %[[SPMAT]], %[[DNB]]{TRANSPOSE}, %[[DNC]] : index into f64
+// CHECK: gpu.spmm async {{\[.*\]}} %[[SPMAT]], %[[DNB]]{TRANSPOSE}, %[[DNC]],
+// CHECK: return
+// CHECK: }
+func.func @matmul_transpose_b(%A: tensor<?x?xf64, #SortedCOO>,
+ %B: tensor<?x?xf64>,
+ %C_in: tensor<?x?xf64>) -> tensor<?x?xf64> {
+ %C_out = linalg.matmul
+ indexing_maps = [
+ affine_map<(i, j, k) -> (i, k)>,
+ affine_map<(i, j, k) -> (j, k)>,
+ affine_map<(i, j, k) -> (i, j)>
+ ]
+ ins(%A, %B: tensor<?x?xf64, #SortedCOO>, tensor<?x?xf64>)
+ outs(%C_in: tensor<?x?xf64>) -> tensor<?x?xf64>
+ return %C_out : tensor<?x?xf64>
+}
+
+// CHECK-LABEL: func.func @matmul_transpose_ab(
+// CHECK-SAME: %[[A:.*]]: tensor<?x?xf64, #sparse{{[0-9]*}}>,
+// CHECK-SAME: %[[B:.*]]: tensor<?x?xf64>,
+// CHECK-SAME: %[[C:.*]]: tensor<?x?xf64>) -> tensor<?x?xf64> {
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[NNZ:.*]] = sparse_tensor.number_of_entries %[[A]]
+// CHECK-DAG: %[[ROWS:.*]] = tensor.dim %[[A]], %[[C0]]
+// CHECK-DAG: %[[COLS:.*]] = tensor.dim %[[A]], %[[C1]]
+// CHECK-DAG: %[[N:.*]] = tensor.dim %[[B]], %[[C0]]
+// CHECK-DAG: %[[B_HOST:.*]] = bufferization.to_buffer %[[B]]
+// CHECK-DAG: gpu.memcpy async {{\[.*\]}} %[[B_DEVICE:.*]], %[[B_HOST]] :
+// CHECK-DAG: %[[C_HOST:.*]] = bufferization.to_buffer %[[C]]
+// CHECK-DAG: gpu.memcpy async {{\[.*\]}} %[[C_DEVICE:.*]], %[[C_HOST]] :
+// CHECK: %[[SPMAT:.*]], {{.*}} = gpu.create_coo async {{\[.*\]}} %[[ROWS]], %[[COLS]], %[[NNZ]],
+// CHECK: %[[DNB:.*]], {{.*}} = gpu.create_dn_tensor async {{\[.*\]}} %[[B_DEVICE]], %[[N]], %[[ROWS]] : index, index into memref<?x?xf64>
+// CHECK: %[[DNC:.*]], {{.*}} = gpu.create_dn_tensor async {{\[.*\]}} %[[C_DEVICE]], %[[COLS]], %[[N]] : index, index into memref<?x?xf64>
+// CHECK: gpu.spmm_buffer_size async {{\[.*\]}} %[[SPMAT]]{TRANSPOSE}, %[[DNB]]{TRANSPOSE}, %[[DNC]] : index into f64
+// CHECK: gpu.spmm async {{\[.*\]}} %[[SPMAT]]{TRANSPOSE}, %[[DNB]]{TRANSPOSE}, %[[DNC]],
+// CHECK: return
+// CHECK: }
+func.func @matmul_transpose_ab(%A: tensor<?x?xf64, #SortedCOO>,
+ %B: tensor<?x?xf64>,
+ %C_in: tensor<?x?xf64>) -> tensor<?x?xf64> {
+ %C_out = linalg.matmul
+ indexing_maps = [
+ affine_map<(i, j, k) -> (k, i)>,
+ affine_map<(i, j, k) -> (j, k)>,
+ affine_map<(i, j, k) -> (i, j)>
+ ]
+ ins(%A, %B: tensor<?x?xf64, #SortedCOO>, tensor<?x?xf64>)
+ outs(%C_in: tensor<?x?xf64>) -> tensor<?x?xf64>
+ return %C_out : tensor<?x?xf64>
+}
+
+#CSR = #sparse_tensor.encoding<{
+ map = (d0, d1) -> (d0 : dense, d1 : compressed)
+}>
+
+// CHECK-LABEL: func.func @spgemm_transpose_a(
+// CHECK-NOT: gpu.
+// CHECK: linalg.generic
+// CHECK-NOT: gpu.
+// CHECK: return
+func.func @spgemm_transpose_a(%A: tensor<?x?xf64, #CSR>,
+ %B: tensor<?x?xf64, #CSR>,
+ %C_in: tensor<?x?xf64, #CSR>) -> tensor<?x?xf64, #CSR> {
+ %C_out = linalg.matmul
+ indexing_maps = [
+ affine_map<(i, j, k) -> (k, i)>,
+ affine_map<(i, j, k) -> (k, j)>,
+ affine_map<(i, j, k) -> (i, j)>
+ ]
+ ins(%A, %B: tensor<?x?xf64, #CSR>, tensor<?x?xf64, #CSR>)
+ outs(%C_in: tensor<?x?xf64, #CSR>) -> tensor<?x?xf64, #CSR>
+ return %C_out : tensor<?x?xf64, #CSR>
+}
+
+#NV_24 = #sparse_tensor.encoding<{
+ map = (i, j) -> (i : dense, j floordiv 4 : dense, j mod 4 : structured[2, 4])
+}>
+
+// CHECK-LABEL: func.func @matmul_2to4_transpose_b(
+// CHECK-NOT: gpu.
+// CHECK: linalg.generic
+// CHECK-NOT: gpu.
+// CHECK: return
+func.func @matmul_2to4_transpose_b(%A: tensor<?x?xf16>,
+ %B: tensor<?x?xf16>,
+ %C_in: tensor<?x?xf16>) -> tensor<?x?xf16> {
+ %A_sparse = sparse_tensor.convert %A : tensor<?x?xf16> to tensor<?x?xf16, #NV_24>
+ %C_out = linalg.matmul
+ indexing_maps = [
+ affine_map<(i, j, k) -> (i, k)>,
+ affine_map<(i, j, k) -> (j, k)>,
+ affine_map<(i, j, k) -> (i, j)>
+ ]
+ ins(%A_sparse, %B: tensor<?x?xf16, #NV_24>, tensor<?x?xf16>)
+ outs(%C_in: tensor<?x?xf16>) -> tensor<?x?xf16>
+ return %C_out : tensor<?x?xf16>
+}
diff --git a/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-matmul-transpose-lib.mlir b/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-matmul-transpose-lib.mlir
new file mode 100644
index 00000000000000..b31643d911bec9
--- /dev/null
+++ b/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-matmul-transpose-lib.mlir
@@ -0,0 +1,456 @@
+// NOTE: this test requires gpu-sm80
+//
+// DEFINE: %{compile} = mlir-opt %s \
+// DEFINE: --sparsifier="enable-gpu-libgen gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71 gpu-format=%gpu_compilation_format
+// DEFINE: %{run} = mlir-runner \
+// DEFINE: --shared-libs=%mlir_cuda_runtime \
+// DEFINE: --shared-libs=%mlir_c_runner_utils \
+// DEFINE: --e main --entry-point-result=void \
+// DEFINE: | FileCheck %s
+//
+// with RT lib (SoA COO):
+//
+// RUN: %{compile} enable-runtime-library=true" | %{run}
+//
+// without RT lib (AoS COO): note, may fall back to CPU
+//
+// RUN: %{compile} enable-runtime-library=false" | %{run}
+
+#SortedCOO = #sparse_tensor.encoding<{
+ map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton)
+}>
+
+#CSR = #sparse_tensor.encoding<{
+ map = (d0, d1) -> (d0 : dense, d1 : compressed),
+ posWidth = 32,
+ crdWidth = 32
+}>
+
+#CSC = #sparse_tensor.encoding<{
+ map = (d0, d1) -> (d1 : dense, d0 : compressed),
+ posWidth = 64,
+ crdWidth = 64
+}>
+
+#transpose_a = [
+ affine_map<(i, j, k) -> (k, i)>,
+ affine_map<(i, j, k) -> (k, j)>,
+ affine_map<(i, j, k) -> (i, j)>
+]
+
+#transpose_b = [
+ affine_map<(i, j, k) -> (i, k)>,
+ affine_map<(i, j, k) -> (j, k)>,
+ affine_map<(i, j, k) -> (i, j)>
+]
+
+#transpose_ab = [
+ affine_map<(i, j, k) -> (k, i)>,
+ affine_map<(i, j, k) -> (j, k)>,
+ affine_map<(i, j, k) -> (i, j)>
+]
+
+module {
+ llvm.func @mgpuCreateSparseEnv()
+ llvm.func @mgpuDestroySparseEnv()
+
+ // Compute C = A^T B with A in COO format.
+ func.func @matmulCOO_transpose_a(%A: tensor<4x8xf32, #SortedCOO>,
+ %B: tensor<4x8xf32>,
+ %C: tensor<8x8xf32>) -> tensor<8x8xf32> {
+ %D = linalg.matmul indexing_maps = #transpose_a
+ ins(%A, %B: tensor<4x8xf32, #SortedCOO>, tensor<4x8xf32>)
+ outs(%C: tensor<8x8xf32>) -> tensor<8x8xf32>
+ return %D : tensor<8x8xf32>
+ }
+
+ // Compute C = A^T B with A in CSR format.
+ func.func @matmulCSR_transpose_a(%A: tensor<4x8xf32, #CSR>,
+ %B: tensor<4x8xf32>,
+ %C: tensor<8x8xf32>) -> tensor<8x8xf32> {
+ %D = linalg.matmul indexing_maps = #transpose_a
+ ins(%A, %B: tensor<4x8xf32, #CSR>, tensor<4x8xf32>)
+ outs(%C: tensor<8x8xf32>) -> tensor<8x8xf32>
+ return %D : tensor<8x8xf32>
+ }
+
+ // Compute C = A^T B with A in CSC format.
+ func.func @matmulCSC_transpose_a(%A: tensor<4x8xf32, #CSC>,
+ %B: tensor<4x8xf32>,
+ %C: tensor<8x8xf32>) -> tensor<8x8xf32> {
+ %D = linalg.matmul indexing_maps = #transpose_a
+ ins(%A, %B: tensor<4x8xf32, #CSC>, tensor<4x8xf32>)
+ outs(%C: tensor<8x8xf32>) -> tensor<8x8xf32>
+ return %D : tensor<8x8xf32>
+ }
+
+ // Compute C = A B^T with A in COO format.
+ func.func @matmulCOO_transpose_b(%A: tensor<8x4xf32, #SortedCOO>,
+ %B: tensor<8x4xf32>,
+ %C: tensor<8x8xf32>) -> tensor<8x8xf32> {
+ %D = linalg.matmul indexing_maps = #transpose_b
+ ins(%A, %B: tensor<8x4xf32, #SortedCOO>, tensor<8x4xf32>)
+ outs(%C: tensor<8x8xf32>) -> tensor<8x8xf32>
+ return %D : tensor<8x8xf32>
+ }
+
+ // Compute C = A B^T with A in CSR format.
+ func.func @matmulCSR_transpose_b(%A: tensor<8x4xf32, #CSR>,
+ %B: tensor<8x4xf32>,
+ %C: tensor<8x8xf32>) -> tensor<8x8xf32> {
+ %D = linalg.matmul indexing_maps = #transpose_b
+ ins(%A, %B: tensor<8x4xf32, #CSR>, tensor<8x4xf32>)
+ outs(%C: tensor<8x8xf32>) -> tensor<8x8xf32>
+ return %D : tensor<8x8xf32>
+ }
+
+ // Compute C = A B^T with A in CSC format.
+ func.func @matmulCSC_transpose_b(%A: tensor<8x4xf32, #CSC>,
+ %B: tensor<8x4xf32>,
+ %C: tensor<8x8xf32>) -> tensor<8x8xf32> {
+ %D = linalg.matmul indexing_maps = #transpose_b
+ ins(%A, %B: tensor<8x4xf32, #CSC>, tensor<8x4xf32>)
+ outs(%C: tensor<8x8xf32>) -> tensor<8x8xf32>
+ return %D : tensor<8x8xf32>
+ }
+
+ // Compute C = A^T B^T with A in COO format.
+ func.func @matmulCOO_transpose_ab(%A: tensor<4x8xf32, #SortedCOO>,
+ %B: tensor<8x4xf32>,
+ %C: tensor<8x8xf32>) -> tensor<8x8xf32> {
+ %D = linalg.matmul indexing_maps = #transpose_ab
+ ins(%A, %B: tensor<4x8xf32, #SortedCOO>, tensor<8x4xf32>)
+ outs(%C: tensor<8x8xf32>) -> tensor<8x8xf32>
+ return %D : tensor<8x8xf32>
+ }
+
+ // Compute C = A^T B^T with A in CSR format.
+ func.func @matmulCSR_transpose_ab(%A: tensor<4x8xf32, #CSR>,
+ %B: tensor<8x4xf32>,
+ %C: tensor<8x8xf32>) -> tensor<8x8xf32> {
+ %D = linalg.matmul indexing_maps = #transpose_ab
+ ins(%A, %B: tensor<4x8xf32, #CSR>, tensor<8x4xf32>)
+ outs(%C: tensor<8x8xf32>) -> tensor<8x8xf32>
+ return %D : tensor<8x8xf32>
+ }
+
+ // Compute C = A^T B^T with A in CSC format.
+ func.func @matmulCSC_transpose_ab(%A: tensor<4x8xf32, #CSC>,
+ %B: tensor<8x4xf32>,
+ %C: tensor<8x8xf32>) -> tensor<8x8xf32> {
+ %D = linalg.matmul indexing_maps = #transpose_ab
+ ins(%A, %B: tensor<4x8xf32, #CSC>, tensor<8x4xf32>)
+ outs(%C: tensor<8x8xf32>) -> tensor<8x8xf32>
+ return %D : tensor<8x8xf32>
+ }
+
+ // Helper to dump dense tensor as series of vectors.
+ func.func @dump(%mat: tensor<8x8xf32>) {
+ %f0 = arith.constant 0.0 : f32
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c8 = arith.constant 8 : index
+ scf.for %i = %c0 to %c8 step %c1 {
+ %v = vector.transfer_read %mat[%i,%c0], %f0 : tensor<8x8xf32>, vector<8xf32>
+ vector.print %v : vector<8xf32>
+ }
+ return
+ }
+
+ //
+ // Main driver.
+ //
+ func.func @main() {
+ llvm.call @mgpuCreateSparseEnv() : () -> ()
+ %f0 = arith.constant 0.0 : f32
+ %f1 = arith.constant 1.0 : f32
+
+ // Stress test with rectangular dense matrices DA and DAT.
+ %DA = tensor.generate {
+ ^bb0(%i: index, %j: index):
+ %k = arith.addi %i, %j : index
+ %l = arith.index_cast %k : index to i64
+ %f = arith.uitofp %l : i64 to f32
+ tensor.yield %f : f32
+ } : tensor<8x4xf32>
+ %DAT = tensor.generate {
+ ^bb0(%i: index, %j: index):
+ %k = arith.addi %i, %j : index
+ %l = arith.index_cast %k : index to i64
+ %f = arith.uitofp %l : i64 to f32
+ tensor.yield %f : f32
+ } : tensor<4x8xf32>
+
+ // Convert to "sparse" matrices A and AT.
+ %Acoo = sparse_tensor.convert %DA : tensor<8x4xf32> to tensor<8x4xf32, #SortedCOO>
+ %Acsr = sparse_tensor.convert %DA : tensor<8x4xf32> to tensor<8x4xf32, #CSR>
+ %Acsc = sparse_tensor.convert %DA : tensor<8x4xf32> to tensor<8x4xf32, #CSC>
+ %ATcoo = sparse_tensor.convert %DAT : tensor<4x8xf32> to tensor<4x8xf32, #SortedCOO>
+ %ATcsr = sparse_tensor.convert %DAT : tensor<4x8xf32> to tensor<4x8xf32, #CSR>
+ %ATcsc = sparse_tensor.convert %DAT : tensor<4x8xf32> to tensor<4x8xf32, #CSC>
+
+ // Initial C matrices.
+ %C0 = tensor.generate {
+ ^bb0(%i: index, %j: index):
+ tensor.yield %f0 : f32
+ } : tensor<8x8xf32>
+ %C1 = tensor.generate {
+ ^bb0(%i: index, %j: index):
+ tensor.yield %f1 : f32
+ } : tensor<8x8xf32>
+
+ // Call the kernels.
+ %0 = call @matmulCOO_transpose_a(%ATcoo, %DAT, %C0) : (tensor<4x8xf32, #SortedCOO>,
+ tensor<4x8xf32>,
+ tensor<8x8xf32>) -> tensor<8x8xf32>
+ %1 = call @matmulCSR_transpose_a(%ATcsr, %DAT, %C0) : (tensor<4x8xf32, #CSR>,
+ tensor<4x8xf32>,
+ tensor<8x8xf32>) -> tensor<8x8xf32>
+ %2 = call @matmulCSC_transpose_a(%ATcsc, %DAT, %C0) : (tensor<4x8xf32, #CSC>,
+ tensor<4x8xf32>,
+ tensor<8x8xf32>) -> tensor<8x8xf32>
+ %3 = call @matmulCOO_transpose_a(%ATcoo, %DAT, %C1) : (tensor<4x8xf32, #SortedCOO>,
+ tensor<4x8xf32>,
+ tensor<8x8xf32>) -> tensor<8x8xf32>
+ %4 = call @matmulCSR_transpose_a(%ATcsr, %DAT, %C1) : (tensor<4x8xf32, #CSR>,
+ tensor<4x8xf32>,
+ tensor<8x8xf32>) -> tensor<8x8xf32>
+ %5 = call @matmulCSC_transpose_a(%ATcsc, %DAT, %C1) : (tensor<4x8xf32, #CSC>,
+ tensor<4x8xf32>,
+ tensor<8x8xf32>) -> tensor<8x8xf32>
+ %6 = call @matmulCOO_transpose_b(%Acoo, %DA, %C0) : (tensor<8x4xf32, #SortedCOO>,
+ tensor<8x4xf32>,
+ tensor<8x8xf32>) -> tensor<8x8xf32>
+ %7 = call @matmulCSR_transpose_b(%Acsr, %DA, %C0) : (tensor<8x4xf32, #CSR>,
+ tensor<8x4xf32>,
+ tensor<8x8xf32>) -> tensor<8x8xf32>
+ %8 = call @matmulCSC_transpose_b(%Acsc, %DA, %C0) : (tensor<8x4xf32, #CSC>,
+ tensor<8x4xf32>,
+ tensor<8x8xf32>) -> tensor<8x8xf32>
+ %9 = call @matmulCOO_transpose_b(%Acoo, %DA, %C1) : (tensor<8x4xf32, #SortedCOO>,
+ tensor<8x4xf32>,
+ tensor<8x8xf32>) -> tensor<8x8xf32>
+ %10 = call @matmulCSR_transpose_b(%Acsr, %DA, %C1) : (tensor<8x4xf32, #CSR>,
+ tensor<8x4xf32>,
+ tensor<8x8xf32>) -> tensor<8x8xf32>
+ %11 = call @matmulCSC_transpose_b(%Acsc, %DA, %C1) : (tensor<8x4xf32, #CSC>,
+ tensor<8x4xf32>,
+ tensor<8x8xf32>) -> tensor<8x8xf32>
+ %12 = call @matmulCOO_transpose_ab(%ATcoo, %DA, %C0) : (tensor<4x8xf32, #SortedCOO>,
+ tensor<8x4xf32>,
+ tensor<8x8xf32>) -> tensor<8x8xf32>
+ %13 = call @matmulCSR_transpose_ab(%ATcsr, %DA, %C0) : (tensor<4x8xf32, #CSR>,
+ tensor<8x4xf32>,
+ tensor<8x8xf32>) -> tensor<8x8xf32>
+ %14 = call @matmulCSC_transpose_ab(%ATcsc, %DA, %C0) : (tensor<4x8xf32, #CSC>,
+ tensor<8x4xf32>,
+ tensor<8x8xf32>) -> tensor<8x8xf32>
+ %15 = call @matmulCOO_transpose_ab(%ATcoo, %DA, %C1) : (tensor<4x8xf32, #SortedCOO>,
+ tensor<8x4xf32>,
+ tensor<8x8xf32>) -> tensor<8x8xf32>
+ %16 = call @matmulCSR_transpose_ab(%ATcsr, %DA, %C1) : (tensor<4x8xf32, #CSR>,
+ tensor<8x4xf32>,
+ tensor<8x8xf32>) -> tensor<8x8xf32>
+ %17 = call @matmulCSC_transpose_ab(%ATcsc, %DA, %C1) : (tensor<4x8xf32, #CSC>,
+ tensor<8x4xf32>,
+ tensor<8x8xf32>) -> tensor<8x8xf32>
+
+ //
+ // Sanity check on results.
+ //
+ // CHECK: ( 14, 20, 26, 32, 38, 44, 50, 56 )
+ // CHECK-NEXT: ( 20, 30, 40, 50, 60, 70, 80, 90 )
+ // CHECK-NEXT: ( 26, 40, 54, 68, 82, 96, 110, 124 )
+ // CHECK-NEXT: ( 32, 50, 68, 86, 104, 122, 140, 158 )
+ // CHECK-NEXT: ( 38, 60, 82, 104, 126, 148, 170, 192 )
+ // CHECK-NEXT: ( 44, 70, 96, 122, 148, 174, 200, 226 )
+ // CHECK-NEXT: ( 50, 80, 110, 140, 170, 200, 230, 260 )
+ // CHECK-NEXT: ( 56, 90, 124, 158, 192, 226, 260, 294 )
+ //
+ // CHECK: ( 14, 20, 26, 32, 38, 44, 50, 56 )
+ // CHECK-NEXT: ( 20, 30, 40, 50, 60, 70, 80, 90 )
+ // CHECK-NEXT: ( 26, 40, 54, 68, 82, 96, 110, 124 )
+ // CHECK-NEXT: ( 32, 50, 68, 86, 104, 122, 140, 158 )
+ // CHECK-NEXT: ( 38, 60, 82, 104, 126, 148, 170, 192 )
+ // CHECK-NEXT: ( 44, 70, 96, 122, 148, 174, 200, 226 )
+ // CHECK-NEXT: ( 50, 80, 110, 140, 170, 200, 230, 260 )
+ // CHECK-NEXT: ( 56, 90, 124, 158, 192, 226, 260, 294 )
+ //
+ // CHECK: ( 14, 20, 26, 32, 38, 44, 50, 56 )
+ // CHECK-NEXT: ( 20, 30, 40, 50, 60, 70, 80, 90 )
+ // CHECK-NEXT: ( 26, 40, 54, 68, 82, 96, 110, 124 )
+ // CHECK-NEXT: ( 32, 50, 68, 86, 104, 122, 140, 158 )
+ // CHECK-NEXT: ( 38, 60, 82, 104, 126, 148, 170, 192 )
+ // CHECK-NEXT: ( 44, 70, 96, 122, 148, 174, 200, 226 )
+ // CHECK-NEXT: ( 50, 80, 110, 140, 170, 200, 230, 260 )
+ // CHECK-NEXT: ( 56, 90, 124, 158, 192, 226, 260, 294 )
+ //
+ // CHECK: ( 15, 21, 27, 33, 39, 45, 51, 57 )
+ // CHECK-NEXT: ( 21, 31, 41, 51, 61, 71, 81, 91 )
+ // CHECK-NEXT: ( 27, 41, 55, 69, 83, 97, 111, 125 )
+ // CHECK-NEXT: ( 33, 51, 69, 87, 105, 123, 141, 159 )
+ // CHECK-NEXT: ( 39, 61, 83, 105, 127, 149, 171, 193 )
+ // CHECK-NEXT: ( 45, 71, 97, 123, 149, 175, 201, 227 )
+ // CHECK-NEXT: ( 51, 81, 111, 141, 171, 201, 231, 261 )
+ // CHECK-NEXT: ( 57, 91, 125, 159, 193, 227, 261, 295 )
+ //
+ // CHECK: ( 15, 21, 27, 33, 39, 45, 51, 57 )
+ // CHECK-NEXT: ( 21, 31, 41, 51, 61, 71, 81, 91 )
+ // CHECK-NEXT: ( 27, 41, 55, 69, 83, 97, 111, 125 )
+ // CHECK-NEXT: ( 33, 51, 69, 87, 105, 123, 141, 159 )
+ // CHECK-NEXT: ( 39, 61, 83, 105, 127, 149, 171, 193 )
+ // CHECK-NEXT: ( 45, 71, 97, 123, 149, 175, 201, 227 )
+ // CHECK-NEXT: ( 51, 81, 111, 141, 171, 201, 231, 261 )
+ // CHECK-NEXT: ( 57, 91, 125, 159, 193, 227, 261, 295 )
+ //
+ // CHECK: ( 15, 21, 27, 33, 39, 45, 51, 57 )
+ // CHECK-NEXT: ( 21, 31, 41, 51, 61, 71, 81, 91 )
+ // CHECK-NEXT: ( 27, 41, 55, 69, 83, 97, 111, 125 )
+ // CHECK-NEXT: ( 33, 51, 69, 87, 105, 123, 141, 159 )
+ // CHECK-NEXT: ( 39, 61, 83, 105, 127, 149, 171, 193 )
+ // CHECK-NEXT: ( 45, 71, 97, 123, 149, 175, 201, 227 )
+ // CHECK-NEXT: ( 51, 81, 111, 141, 171, 201, 231, 261 )
+ // CHECK-NEXT: ( 57, 91, 125, 159, 193, 227, 261, 295 )
+ //
+ // A B^T with C initialized to zeros, then ones.
+ // CHECK: ( 14, 20, 26, 32, 38, 44, 50, 56 )
+ // CHECK-NEXT: ( 20, 30, 40, 50, 60, 70, 80, 90 )
+ // CHECK-NEXT: ( 26, 40, 54, 68, 82, 96, 110, 124 )
+ // CHECK-NEXT: ( 32, 50, 68, 86, 104, 122, 140, 158 )
+ // CHECK-NEXT: ( 38, 60, 82, 104, 126, 148, 170, 192 )
+ // CHECK-NEXT: ( 44, 70, 96, 122, 148, 174, 200, 226 )
+ // CHECK-NEXT: ( 50, 80, 110, 140, 170, 200, 230, 260 )
+ // CHECK-NEXT: ( 56, 90, 124, 158, 192, 226, 260, 294 )
+ //
+ // CHECK: ( 14, 20, 26, 32, 38, 44, 50, 56 )
+ // CHECK-NEXT: ( 20, 30, 40, 50, 60, 70, 80, 90 )
+ // CHECK-NEXT: ( 26, 40, 54, 68, 82, 96, 110, 124 )
+ // CHECK-NEXT: ( 32, 50, 68, 86, 104, 122, 140, 158 )
+ // CHECK-NEXT: ( 38, 60, 82, 104, 126, 148, 170, 192 )
+ // CHECK-NEXT: ( 44, 70, 96, 122, 148, 174, 200, 226 )
+ // CHECK-NEXT: ( 50, 80, 110, 140, 170, 200, 230, 260 )
+ // CHECK-NEXT: ( 56, 90, 124, 158, 192, 226, 260, 294 )
+ //
+ // CHECK: ( 14, 20, 26, 32, 38, 44, 50, 56 )
+ // CHECK-NEXT: ( 20, 30, 40, 50, 60, 70, 80, 90 )
+ // CHECK-NEXT: ( 26, 40, 54, 68, 82, 96, 110, 124 )
+ // CHECK-NEXT: ( 32, 50, 68, 86, 104, 122, 140, 158 )
+ // CHECK-NEXT: ( 38, 60, 82, 104, 126, 148, 170, 192 )
+ // CHECK-NEXT: ( 44, 70, 96, 122, 148, 174, 200, 226 )
+ // CHECK-NEXT: ( 50, 80, 110, 140, 170, 200, 230, 260 )
+ // CHECK-NEXT: ( 56, 90, 124, 158, 192, 226, 260, 294 )
+ //
+ // CHECK: ( 15, 21, 27, 33, 39, 45, 51, 57 )
+ // CHECK-NEXT: ( 21, 31, 41, 51, 61, 71, 81, 91 )
+ // CHECK-NEXT: ( 27, 41, 55, 69, 83, 97, 111, 125 )
+ // CHECK-NEXT: ( 33, 51, 69, 87, 105, 123, 141, 159 )
+ // CHECK-NEXT: ( 39, 61, 83, 105, 127, 149, 171, 193 )
+ // CHECK-NEXT: ( 45, 71, 97, 123, 149, 175, 201, 227 )
+ // CHECK-NEXT: ( 51, 81, 111, 141, 171, 201, 231, 261 )
+ // CHECK-NEXT: ( 57, 91, 125, 159, 193, 227, 261, 295 )
+ //
+ // CHECK: ( 15, 21, 27, 33, 39, 45, 51, 57 )
+ // CHECK-NEXT: ( 21, 31, 41, 51, 61, 71, 81, 91 )
+ // CHECK-NEXT: ( 27, 41, 55, 69, 83, 97, 111, 125 )
+ // CHECK-NEXT: ( 33, 51, 69, 87, 105, 123, 141, 159 )
+ // CHECK-NEXT: ( 39, 61, 83, 105, 127, 149, 171, 193 )
+ // CHECK-NEXT: ( 45, 71, 97, 123, 149, 175, 201, 227 )
+ // CHECK-NEXT: ( 51, 81, 111, 141, 171, 201, 231, 261 )
+ // CHECK-NEXT: ( 57, 91, 125, 159, 193, 227, 261, 295 )
+ //
+ // CHECK: ( 15, 21, 27, 33, 39, 45, 51, 57 )
+ // CHECK-NEXT: ( 21, 31, 41, 51, 61, 71, 81, 91 )
+ // CHECK-NEXT: ( 27, 41, 55, 69, 83, 97, 111, 125 )
+ // CHECK-NEXT: ( 33, 51, 69, 87, 105, 123, 141, 159 )
+ // CHECK-NEXT: ( 39, 61, 83, 105, 127, 149, 171, 193 )
+ // CHECK-NEXT: ( 45, 71, 97, 123, 149, 175, 201, 227 )
+ // CHECK-NEXT: ( 51, 81, 111, 141, 171, 201, 231, 261 )
+ // CHECK-NEXT: ( 57, 91, 125, 159, 193, 227, 261, 295 )
+ //
+ // A^T B^T with C initialized to zeros, then ones.
+ // CHECK: ( 14, 20, 26, 32, 38, 44, 50, 56 )
+ // CHECK-NEXT: ( 20, 30, 40, 50, 60, 70, 80, 90 )
+ // CHECK-NEXT: ( 26, 40, 54, 68, 82, 96, 110, 124 )
+ // CHECK-NEXT: ( 32, 50, 68, 86, 104, 122, 140, 158 )
+ // CHECK-NEXT: ( 38, 60, 82, 104, 126, 148, 170, 192 )
+ // CHECK-NEXT: ( 44, 70, 96, 122, 148, 174, 200, 226 )
+ // CHECK-NEXT: ( 50, 80, 110, 140, 170, 200, 230, 260 )
+ // CHECK-NEXT: ( 56, 90, 124, 158, 192, 226, 260, 294 )
+ //
+ // CHECK: ( 14, 20, 26, 32, 38, 44, 50, 56 )
+ // CHECK-NEXT: ( 20, 30, 40, 50, 60, 70, 80, 90 )
+ // CHECK-NEXT: ( 26, 40, 54, 68, 82, 96, 110, 124 )
+ // CHECK-NEXT: ( 32, 50, 68, 86, 104, 122, 140, 158 )
+ // CHECK-NEXT: ( 38, 60, 82, 104, 126, 148, 170, 192 )
+ // CHECK-NEXT: ( 44, 70, 96, 122, 148, 174, 200, 226 )
+ // CHECK-NEXT: ( 50, 80, 110, 140, 170, 200, 230, 260 )
+ // CHECK-NEXT: ( 56, 90, 124, 158, 192, 226, 260, 294 )
+ //
+ // CHECK: ( 14, 20, 26, 32, 38, 44, 50, 56 )
+ // CHECK-NEXT: ( 20, 30, 40, 50, 60, 70, 80, 90 )
+ // CHECK-NEXT: ( 26, 40, 54, 68, 82, 96, 110, 124 )
+ // CHECK-NEXT: ( 32, 50, 68, 86, 104, 122, 140, 158 )
+ // CHECK-NEXT: ( 38, 60, 82, 104, 126, 148, 170, 192 )
+ // CHECK-NEXT: ( 44, 70, 96, 122, 148, 174, 200, 226 )
+ // CHECK-NEXT: ( 50, 80, 110, 140, 170, 200, 230, 260 )
+ // CHECK-NEXT: ( 56, 90, 124, 158, 192, 226, 260, 294 )
+ //
+ // CHECK: ( 15, 21, 27, 33, 39, 45, 51, 57 )
+ // CHECK-NEXT: ( 21, 31, 41, 51, 61, 71, 81, 91 )
+ // CHECK-NEXT: ( 27, 41, 55, 69, 83, 97, 111, 125 )
+ // CHECK-NEXT: ( 33, 51, 69, 87, 105, 123, 141, 159 )
+ // CHECK-NEXT: ( 39, 61, 83, 105, 127, 149, 171, 193 )
+ // CHECK-NEXT: ( 45, 71, 97, 123, 149, 175, 201, 227 )
+ // CHECK-NEXT: ( 51, 81, 111, 141, 171, 201, 231, 261 )
+ // CHECK-NEXT: ( 57, 91, 125, 159, 193, 227, 261, 295 )
+ //
+ // CHECK: ( 15, 21, 27, 33, 39, 45, 51, 57 )
+ // CHECK-NEXT: ( 21, 31, 41, 51, 61, 71, 81, 91 )
+ // CHECK-NEXT: ( 27, 41, 55, 69, 83, 97, 111, 125 )
+ // CHECK-NEXT: ( 33, 51, 69, 87, 105, 123, 141, 159 )
+ // CHECK-NEXT: ( 39, 61, 83, 105, 127, 149, 171, 193 )
+ // CHECK-NEXT: ( 45, 71, 97, 123, 149, 175, 201, 227 )
+ // CHECK-NEXT: ( 51, 81, 111, 141, 171, 201, 231, 261 )
+ // CHECK-NEXT: ( 57, 91, 125, 159, 193, 227, 261, 295 )
+ //
+ // CHECK: ( 15, 21, 27, 33, 39, 45, 51, 57 )
+ // CHECK-NEXT: ( 21, 31, 41, 51, 61, 71, 81, 91 )
+ // CHECK-NEXT: ( 27, 41, 55, 69, 83, 97, 111, 125 )
+ // CHECK-NEXT: ( 33, 51, 69, 87, 105, 123, 141, 159 )
+ // CHECK-NEXT: ( 39, 61, 83, 105, 127, 149, 171, 193 )
+ // CHECK-NEXT: ( 45, 71, 97, 123, 149, 175, 201, 227 )
+ // CHECK-NEXT: ( 51, 81, 111, 141, 171, 201, 231, 261 )
+ // CHECK-NEXT: ( 57, 91, 125, 159, 193, 227, 261, 295 )
+ //
+ call @dump(%0) : (tensor<8x8xf32>) -> ()
+ call @dump(%1) : (tensor<8x8xf32>) -> ()
+ call @dump(%2) : (tensor<8x8xf32>) -> ()
+ call @dump(%3) : (tensor<8x8xf32>) -> ()
+ call @dump(%4) : (tensor<8x8xf32>) -> ()
+ call @dump(%5) : (tensor<8x8xf32>) -> ()
+ call @dump(%6) : (tensor<8x8xf32>) -> ()
+ call @dump(%7) : (tensor<8x8xf32>) -> ()
+ call @dump(%8) : (tensor<8x8xf32>) -> ()
+ call @dump(%9) : (tensor<8x8xf32>) -> ()
+ call @dump(%10) : (tensor<8x8xf32>) -> ()
+ call @dump(%11) : (tensor<8x8xf32>) -> ()
+ call @dump(%12) : (tensor<8x8xf32>) -> ()
+ call @dump(%13) : (tensor<8x8xf32>) -> ()
+ call @dump(%14) : (tensor<8x8xf32>) -> ()
+ call @dump(%15) : (tensor<8x8xf32>) -> ()
+ call @dump(%16) : (tensor<8x8xf32>) -> ()
+ call @dump(%17) : (tensor<8x8xf32>) -> ()
+
+ // Release the resources.
+ bufferization.dealloc_tensor %Acoo : tensor<8x4xf32, #SortedCOO>
+ bufferization.dealloc_tensor %Acsr : tensor<8x4xf32, #CSR>
+ bufferization.dealloc_tensor %Acsc : tensor<8x4xf32, #CSC>
+ bufferization.dealloc_tensor %ATcoo : tensor<4x8xf32, #SortedCOO>
+ bufferization.dealloc_tensor %ATcsr : tensor<4x8xf32, #CSR>
+ bufferization.dealloc_tensor %ATcsc : tensor<4x8xf32, #CSC>
+
+ llvm.call @mgpuDestroySparseEnv() : () -> ()
+
+ return
+ }
+}
More information about the Mlir-commits
mailing list