1// NOTE: this test requires gpu-sm80 2// 3// DEFINE: %{compile} = mlir-opt %s \ 4// DEFINE: --sparsifier="enable-gpu-libgen gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71 gpu-format=%gpu_compilation_format 5// DEFINE: %{run} = \ 6// DEFINE: env TENSOR0="%mlir_src_dir/test/Integration/data/block.mtx" \ 7// DEFINE: mlir-runner \ 8// DEFINE: --shared-libs=%mlir_cuda_runtime \ 9// DEFINE: --shared-libs=%mlir_c_runner_utils \ 10// DEFINE: --e main --entry-point-result=void \ 11// DEFINE: | FileCheck %s 12// 13// with RT lib: 14// 15// RUN: %{compile} enable-runtime-library=true" | %{run} 16// 17// without RT lib: 18// 19// RUN: %{compile} enable-runtime-library=false" | %{run} 20 21!Filename = !llvm.ptr 22 23#CSR = #sparse_tensor.encoding<{ 24 map = (d0, d1) -> (d0 : dense, d1 : compressed) 25}> 26 27#BSR = #sparse_tensor.encoding<{ 28 map = (i, j) -> ( 29 i floordiv 2 : dense, 30 j floordiv 2 : compressed, 31 i mod 2 : dense, 32 j mod 2 : dense) 33}> 34 35#trait_SDDMM = { 36 indexing_maps = [ 37 affine_map<(i,j,k) -> (i,k)>, // A 38 affine_map<(i,j,k) -> (k,j)>, // B 39 affine_map<(i,j,k) -> (i,j)> // S (in/out) 40 ], 41 iterator_types = ["parallel", "parallel", "reduction"], 42 doc = "S(i,j) += spy[S(i,j)] x SUM_k A(i,k) B(k,j)" 43} 44 45// 46// Integration test that lowers a kernel annotated as sparse to 47// actual sparse code, initializes sparse storage schemes, and 48// runs the resulting code with the JIT compiler. 49// 50module { 51 llvm.func @mgpuCreateSparseEnv() 52 llvm.func @mgpuDestroySparseEnv() 53 54 // 55 // A kernel that computes a CSR sampled dense matrix matrix multiplication 56 // using a "spy" function and in-place update of the sampling sparse matrix. 57 // 58 func.func @SDDMM(%args: tensor<?x?xf32, #CSR>, 59 %arga: tensor<?x?xf32>, 60 %argb: tensor<?x?xf32>) -> tensor<?x?xf32, #CSR> { 61 %result = linalg.generic #trait_SDDMM 62 ins(%arga, %argb: tensor<?x?xf32>, tensor<?x?xf32>) 63 outs(%args: tensor<?x?xf32, #CSR>) { 64 ^bb(%a: f32, %b: f32, %s: f32): 65 %f0 = arith.constant 0.0 : f32 66 %u = sparse_tensor.unary %s : f32 to f32 67 present={ 68 ^bb0(%p: f32): 69 %mul = arith.mulf %a, %b : f32 70 sparse_tensor.yield %mul : f32 71 } 72 absent={} 73 %r = sparse_tensor.reduce %s, %u, %f0 : f32 { 74 ^bb0(%p: f32, %q: f32): 75 %add = arith.addf %p, %q : f32 76 sparse_tensor.yield %add : f32 77 } 78 linalg.yield %r : f32 79 } -> tensor<?x?xf32, #CSR> 80 return %result : tensor<?x?xf32, #CSR> 81 } 82 83 // 84 // A kernel that computes a BSR sampled dense matrix matrix multiplication 85 // using a "spy" function and in-place update of the sampling sparse matrix. 86 // 87 func.func @SDDMM_block(%args: tensor<?x?xf32, #BSR>, 88 %arga: tensor<?x?xf32>, 89 %argb: tensor<?x?xf32>) -> tensor<?x?xf32, #BSR> { 90 %result = linalg.generic #trait_SDDMM 91 ins(%arga, %argb: tensor<?x?xf32>, tensor<?x?xf32>) 92 outs(%args: tensor<?x?xf32, #BSR>) { 93 ^bb(%a: f32, %b: f32, %s: f32): 94 %f0 = arith.constant 0.0 : f32 95 %u = sparse_tensor.unary %s : f32 to f32 96 present={ 97 ^bb0(%p: f32): 98 %mul = arith.mulf %a, %b : f32 99 sparse_tensor.yield %mul : f32 100 } 101 absent={} 102 %r = sparse_tensor.reduce %s, %u, %f0 : f32 { 103 ^bb0(%p: f32, %q: f32): 104 %add = arith.addf %p, %q : f32 105 sparse_tensor.yield %add : f32 106 } 107 linalg.yield %r : f32 108 } -> tensor<?x?xf32, #BSR> 109 return %result : tensor<?x?xf32, #BSR> 110 } 111 112 func.func private @getTensorFilename(index) -> (!Filename) 113 114 // 115 // Main driver. 116 // 117 func.func @main() { 118 llvm.call @mgpuCreateSparseEnv() : () -> () 119 %d0 = arith.constant 0.0 : f32 120 %c0 = arith.constant 0 : index 121 %c1 = arith.constant 1 : index 122 %c4 = arith.constant 4 : index 123 %c6 = arith.constant 6 : index 124 125 // Initialize dense matrices. 126 %a = tensor.generate %c4, %c4 { 127 ^bb0(%i: index, %j: index): 128 %p = arith.addi %i, %c1 : index 129 %q = arith.index_cast %p : index to i32 130 %d = arith.sitofp %q : i32 to f32 131 tensor.yield %d : f32 132 } : tensor<?x?xf32> 133 %b = tensor.generate %c4, %c6 { 134 ^bb0(%i: index, %j: index): 135 %p = arith.addi %j, %c1 : index 136 %q = arith.index_cast %p : index to i32 137 %d = arith.sitofp %q : i32 to f32 138 tensor.yield %d : f32 139 } : tensor<?x?xf32> 140 141 // Read the sparse matrix from file, construct sparse storage. 142 // 143 // +-----+-----+-----+ 144 // | 1 2 | . . | 4 . | 145 // | . 3 | . . | . 5 | 146 // +-----+-----+-----+ 147 // | . . | 6 7 | . . | 148 // | . . | 8 . | . . | 149 // +-----+-----+-----+ 150 // 151 %fileName = call @getTensorFilename(%c0) : (index) -> (!Filename) 152 %m_csr = sparse_tensor.new %fileName : !Filename to tensor<?x?xf32, #CSR> 153 %m_bsr = sparse_tensor.new %fileName : !Filename to tensor<?x?xf32, #BSR> 154 155 // Call the kernel. 156 %0 = call @SDDMM(%m_csr, %a, %b) 157 : (tensor<?x?xf32, #CSR>, 158 tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32, #CSR> 159 %1 = call @SDDMM_block(%m_bsr, %a, %b) 160 : (tensor<?x?xf32, #BSR>, 161 tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32, #BSR> 162 163 // 164 // Print the result for verification. Note that the "spy" determines what 165 // dot products are sampled, but the original contents are added back to 166 // the result (which is why the block sparse version has actual results 167 // in the original zero positions). 168 // 169 // CHECK: ---- Sparse Tensor ---- 170 // CHECK-NEXT: nse = 8 171 // CHECK-NEXT: dim = ( 4, 6 ) 172 // CHECK-NEXT: lvl = ( 4, 6 ) 173 // CHECK-NEXT: pos[1] : ( 0, 3, 5, 7, 8 ) 174 // CHECK-NEXT: crd[1] : ( 0, 1, 4, 1, 5, 2, 3, 2 ) 175 // CHECK-NEXT: values : ( 5, 10, 24, 19, 53, 42, 55, 56 ) 176 // CHECK-NEXT: ---- 177 // 178 // CHECK: ---- Sparse Tensor ---- 179 // CHECK-NEXT: nse = 12 180 // CHECK-NEXT: dim = ( 4, 6 ) 181 // CHECK-NEXT: lvl = ( 2, 3, 2, 2 ) 182 // CHECK-NEXT: pos[1] : ( 0, 2, 3 ) 183 // CHECK-NEXT: crd[1] : ( 0, 2, 1 ) 184 // CHECK-NEXT: values : ( 5, 10, 8, 19, 24, 24, 40, 53, 42, 55, 56, 64 ) 185 // CHECK-NEXT: ---- 186 // 187 sparse_tensor.print %0 : tensor<?x?xf32, #CSR> 188 sparse_tensor.print %1 : tensor<?x?xf32, #BSR> 189 190 // Release the resources. 191 bufferization.dealloc_tensor %0 : tensor<?x?xf32, #CSR> 192 bufferization.dealloc_tensor %1 : tensor<?x?xf32, #BSR> 193 194 llvm.call @mgpuDestroySparseEnv() : () -> () 195 return 196 } 197} 198