diff options
| author | Aart Bik <ajcbik@google.com> | 2023-05-11 13:56:34 -0700 |
|---|---|---|
| committer | Aart Bik <ajcbik@google.com> | 2023-05-12 10:44:36 -0700 |
| commit | b700a90cc0aa08252d764b1f7da67bd300469a76 (patch) | |
| tree | ed31487f0e75ac76da26b2ad6f9b8ddd43231a2f /mlir/test/Conversion/GPUCommon | |
| parent | 52761cb99164acd4ea76f91fc16a3e40ec94b898 (diff) | |
| download | llvm-b700a90cc0aa08252d764b1f7da67bd300469a76.zip llvm-b700a90cc0aa08252d764b1f7da67bd300469a76.tar.gz llvm-b700a90cc0aa08252d764b1f7da67bd300469a76.tar.bz2 | |
[mlir][gpu][sparse] add gpu ops for sparse matrix computations
This revision extends the GPU dialect with ops that can be lowered to
host-oriented sparse matrix library calls (in this case cuSparse focused
although the ops could be generalized to support more GPUs in principle).
This will allow the "sparse compiler pipeline" to accelerate sparse operations
(see follow up revisions with examples of this).
For some background;
https://discourse.llvm.org/t/sparse-compiler-and-gpu-code-generation/69786/2
Reviewed By: ThomasRaoux
Differential Revision: https://reviews.llvm.org/D150152
Diffstat (limited to 'mlir/test/Conversion/GPUCommon')
| -rw-r--r-- | mlir/test/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir | 37 |
1 files changed, 37 insertions, 0 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 new file mode 100644 index 0000000..6f163f9 --- /dev/null +++ b/mlir/test/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir @@ -0,0 +1,37 @@ +// RUN: mlir-opt %s --gpu-to-llvm='use-opaque-pointers=1' | FileCheck %s + +module attributes {gpu.container_module} { + + // CHECK-LABEL: func @matvec + // CHECK: llvm.call @mgpuStreamCreate + // CHECK: llvm.call @mgpuMemAlloc + // CHECK: llvm.call @mgpuMemAlloc + // CHECK: llvm.call @mgpuCreateSparseEnv + // CHECK: llvm.call @mgpuCreateCoo + // CHECK: llvm.call @mgpuCreateDnVec + // CHECK: llvm.call @mgpuSpMVBufferSize + // CHECK: llvm.call @mgpuSpM + // CHECK: llvm.call @mgpuDestroySpMat + // CHECK: llvm.call @mgpuDestroyDnVec + // CHECK: llvm.call @mgpuDestroySparseEnv + // CHECK: llvm.call @mgpuStreamSynchronize + // CHECK: llvm.call @mgpuStreamDestroy + func.func @matvec(%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_coo async [%token3] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64> + %dnvec, %token5 = gpu.create_dn_vec async [%token4] %mem2, %arg0 : memref<?xf64> + %bufferSz, %token6 = gpu.spmv_buffer_size async [%token5] %env, %spmat, %dnvec, %dnvec + %token7 = gpu.spmv async [%token6] %env, %spmat, %dnvec, %dnvec, %mem2 : memref<?xf64> + %token8 = gpu.destroy_sp_mat async [%token7] %spmat + %token9 = gpu.destroy_dn_vec async [%token8] %dnvec + %token10 = gpu.destroy_sparse_env async [%token9] %env + gpu.wait [%token10] + return + } + +} + + |