[Mlir-commits] [mlir] [mlir][sparse][gpu] Add libgen support for transposed SpMV and SpMM (PR #230412)
Vito Secona
llvmlistbot at llvm.org
Mon Oct 12 23:50:51 PDT 2026
https://github.com/secona updated https://github.com/llvm/llvm-project/pull/230412
>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/5] [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 d4c0a76088356..2e45813dc1e1c 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 b88492433ba02..c3a1168a1f929 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/5] [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 8243bf9317c22..227f8ea42572a 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 0000000000000..c92a5eb61a4c7
--- /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 0000000000000..c522f17f26bb0
--- /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/5] [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 2e45813dc1e1c..9dcdabe557065 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 c3a1168a1f929..81433f9bcbcc6 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/5] [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 227f8ea42572a..684876597eb95 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 0000000000000..fb333066310a9
--- /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 0000000000000..b31643d911bec
--- /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
+ }
+}
>From cd7c97959ad54536c1e53738f7bec0a2a0fc73a5 Mon Sep 17 00:00:00 2001
From: Vito Secona <secona00 at gmail.com>
Date: Tue, 13 Oct 2026 13:50:13 +0700
Subject: [PATCH 5/5] [mlir][sparse] Handle CONJUGATE_TRANSPOSE
---
.../Transforms/SparseGPUCodegen.cpp | 20 ++++++++++++++++---
1 file changed, 17 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
index 684876597eb95..8a9d8a7307fe3 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
@@ -569,6 +569,18 @@ static Operation *genSpMat(OpBuilder &builder, Location loc,
bCols, bNum, bSz, bSz, rowA, colA, valA);
}
+/// Returns whether the operation swaps matrix dimensions.
+static bool isTransposed(gpu::TransposeMode mode) {
+ switch (mode) {
+ case gpu::TransposeMode::NON_TRANSPOSE:
+ return false;
+ case gpu::TransposeMode::TRANSPOSE:
+ case gpu::TransposeMode::CONJUGATE_TRANSPOSE:
+ return true;
+ }
+ llvm_unreachable("unhandled transpose mode");
+}
+
/// Match and rewrite SpMV kernel.
static LogicalResult rewriteSpMV(PatternRewriter &rewriter,
linalg::GenericOp op, bool enableRT,
@@ -594,7 +606,7 @@ static LogicalResult rewriteSpMV(PatternRewriter &rewriter,
Value nseA = NumberOfEntriesOp::create(rewriter, loc, a);
Value rows = linalg::createOrFoldDimOp(rewriter, loc, a, 0);
Value cols = linalg::createOrFoldDimOp(rewriter, loc, a, 1);
- bool isTranspose = modeA != gpu::TransposeMode::NON_TRANSPOSE;
+ bool isTranspose = isTransposed(modeA);
Value szX = isTranspose ? rows : cols;
Value szY = isTranspose ? cols : rows;
Value memR = genFirstPosOrCrds(rewriter, loc, a, format, enableRT);
@@ -697,8 +709,8 @@ static LogicalResult rewriteSpMM(PatternRewriter &rewriter,
Value nseA = NumberOfEntriesOp::create(rewriter, loc, a);
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;
+ bool isTransposeA = isTransposed(modeA);
+ bool isTransposeB = isTransposed(modeB);
Value szm = isTransposeA ? cols : rows;
Value szk = isTransposeA ? rows : cols;
Value szn = linalg::createOrFoldDimOp(rewriter, loc, b, isTransposeB ? 0 : 1);
@@ -1339,6 +1351,7 @@ struct LinalgOpRewriter : public OpRewritePattern<linalg::GenericOp> {
if (maps == infer({{i, j}, {j}, {i}}))
return rewriteSpMV(rewriter, op, enableRT,
gpu::TransposeMode::NON_TRANSPOSE);
+ // TODO: recognize conjugate transpose patterns for SpMV.
if (maps == infer({{j, i}, {j}, {i}}))
return rewriteSpMV(rewriter, op, enableRT,
gpu::TransposeMode::TRANSPOSE);
@@ -1360,6 +1373,7 @@ struct LinalgOpRewriter : public OpRewritePattern<linalg::GenericOp> {
gpu::TransposeMode::NON_TRANSPOSE,
gpu::TransposeMode::NON_TRANSPOSE);
}
+ // TODO: recognize conjugate transpose patterns for either SpMM input.
if (maps == infer({{k, i}, {k, j}, {i, j}}))
return rewriteSpMM(rewriter, op, enableRT,
gpu::TransposeMode::TRANSPOSE,
More information about the Mlir-commits
mailing list