diff options
| author | Aart Bik <ajcbik@google.com> | 2023-05-15 14:55:45 -0700 |
|---|---|---|
| committer | Aart Bik <ajcbik@google.com> | 2023-05-19 12:46:11 -0700 |
| commit | 981cf1678d6f3399d3fb5e16434669f3992ac6fd (patch) | |
| tree | edbcc013e6661d5c9c16a04f031ee541f40f5bf6 /mlir/test/Conversion/GPUCommon | |
| parent | 94063cac47fb2d07e92d6922761e182f1e657764 (diff) | |
| download | llvm-981cf1678d6f3399d3fb5e16434669f3992ac6fd.zip llvm-981cf1678d6f3399d3fb5e16434669f3992ac6fd.tar.gz llvm-981cf1678d6f3399d3fb5e16434669f3992ac6fd.tar.bz2 | |
[mlir][sparse][gpu] add SpMM to GPU ops dialect
Reviewed By: ThomasRaoux, K-Wu
Differential Revision: https://reviews.llvm.org/D150618
Diffstat (limited to 'mlir/test/Conversion/GPUCommon')
| -rw-r--r-- | mlir/test/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir | 32 |
1 files changed, 31 insertions, 1 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 6f163f9..dcef273 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 @@ -10,7 +10,7 @@ module attributes {gpu.container_module} { // CHECK: llvm.call @mgpuCreateCoo // CHECK: llvm.call @mgpuCreateDnVec // CHECK: llvm.call @mgpuSpMVBufferSize - // CHECK: llvm.call @mgpuSpM + // CHECK: llvm.call @mgpuSpMV // CHECK: llvm.call @mgpuDestroySpMat // CHECK: llvm.call @mgpuDestroyDnVec // CHECK: llvm.call @mgpuDestroySparseEnv @@ -32,6 +32,36 @@ module attributes {gpu.container_module} { return } + // CHECK-LABEL: func @matmul + // CHECK: llvm.call @mgpuStreamCreate + // CHECK: llvm.call @mgpuMemAlloc + // CHECK: llvm.call @mgpuMemAlloc + // CHECK: llvm.call @mgpuCreateSparseEnv + // CHECK: llvm.call @mgpuCreateCsr + // CHECK: llvm.call @mgpuCreateDnMat + // CHECK: llvm.call @mgpuSpMMBufferSize + // CHECK: llvm.call @mgpuSpMM + // CHECK: llvm.call @mgpuDestroySpMat + // CHECK: llvm.call @mgpuDestroyDnMat + // CHECK: llvm.call @mgpuDestroySparseEnv + // CHECK: llvm.call @mgpuStreamSynchronize + // CHECK: llvm.call @mgpuStreamDestroy + func.func @matmul(%arg0: index) { + %token0 = gpu.wait async + %mem1, %token1 = gpu.alloc async [%token0] (%arg0) : memref<?xindex> + %mem2, %token2 = gpu.alloc async [%token1] (%arg0) : memref<?xf64> + %env, %token3 = gpu.create_sparse_env async [%token2] + %spmat, %token4 = gpu.create_csr async [%token3] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64> + %dnmat, %token5 = gpu.create_dn_mat async [%token4] %arg0, %arg0, %mem2 : memref<?xf64> + %bufferSz, %token6 = gpu.spmm_buffer_size async [%token5] %env, %spmat, %dnmat, %dnmat + %token7 = gpu.spmm async [%token6] %env, %spmat, %dnmat, %dnmat, %mem2 : memref<?xf64> + %token8 = gpu.destroy_sp_mat async [%token7] %spmat + %token9 = gpu.destroy_dn_mat async [%token8] %dnmat + %token10 = gpu.destroy_sparse_env async [%token9] %env + gpu.wait [%token10] + return + } + } |