[Mlir-commits] [mlir] [mlir][sparse][gpu] Add libgen support for transposed SpMV and SpMM (PR #230412)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 9 01:48:11 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Vito Secona (secona)
<details>
<summary>Changes</summary>
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
---
Patch is 65.95 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/230412.diff
7 Files Affected:
- (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp (+63-30)
- (modified) mlir/test/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir (+107-2)
- (modified) mlir/test/Dialect/GPU/sparse-roundtrip.mlir (+104)
- (added) mlir/test/Dialect/SparseTensor/GPU/gpu_matmul_transpose_lib.mlir (+157)
- (added) mlir/test/Dialect/SparseTensor/GPU/gpu_matvec_transpose_lib.mlir (+53)
- (added) mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-matmul-transpose-lib.mlir (+456)
- (added) mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-matvec-transpose-lib.mlir (+178)
``````````diff
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
index 8243bf9317c22a..684876597eb95d 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.
@@ -669,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).
@@ -689,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);
@@ -712,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 =
@@ -730,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();
@@ -740,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.
@@ -1324,8 +1335,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.
@@ -1333,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/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir b/mlir/test/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir
index d4c0a76088356f..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
@@ -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
@@ -54,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 b88492433ba02e..81433f9bcbcc6b 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>
@@ -54,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 [%{{.*}}] %{{....
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/230412
More information about the Mlir-commits
mailing list