198 lines · plain
1// NOTE: this test requires gpu-sm802//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_format5// 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 %s12//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.ptr22 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)>, // A38 affine_map<(i,j,k) -> (k,j)>, // B39 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 to47// actual sparse code, initializes sparse storage schemes, and48// 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 multiplication56 // 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_SDDMM62 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 : f3266 %u = sparse_tensor.unary %s : f32 to f3267 present={68 ^bb0(%p: f32):69 %mul = arith.mulf %a, %b : f3270 sparse_tensor.yield %mul : f3271 }72 absent={}73 %r = sparse_tensor.reduce %s, %u, %f0 : f32 {74 ^bb0(%p: f32, %q: f32):75 %add = arith.addf %p, %q : f3276 sparse_tensor.yield %add : f3277 }78 linalg.yield %r : f3279 } -> 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 multiplication85 // 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_SDDMM91 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 : f3295 %u = sparse_tensor.unary %s : f32 to f3296 present={97 ^bb0(%p: f32):98 %mul = arith.mulf %a, %b : f3299 sparse_tensor.yield %mul : f32100 }101 absent={}102 %r = sparse_tensor.reduce %s, %u, %f0 : f32 {103 ^bb0(%p: f32, %q: f32):104 %add = arith.addf %p, %q : f32105 sparse_tensor.yield %add : f32106 }107 linalg.yield %r : f32108 } -> 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 : f32120 %c0 = arith.constant 0 : index121 %c1 = arith.constant 1 : index122 %c4 = arith.constant 4 : index123 %c6 = arith.constant 6 : index124 125 // Initialize dense matrices.126 %a = tensor.generate %c4, %c4 {127 ^bb0(%i: index, %j: index):128 %p = arith.addi %i, %c1 : index129 %q = arith.index_cast %p : index to i32130 %d = arith.sitofp %q : i32 to f32131 tensor.yield %d : f32132 } : tensor<?x?xf32>133 %b = tensor.generate %c4, %c6 {134 ^bb0(%i: index, %j: index):135 %p = arith.addi %j, %c1 : index136 %q = arith.index_cast %p : index to i32137 %d = arith.sitofp %q : i32 to f32138 tensor.yield %d : f32139 } : 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 what165 // dot products are sampled, but the original contents are added back to166 // the result (which is why the block sparse version has actual results167 // in the original zero positions).168 //169 // CHECK: ---- Sparse Tensor ----170 // CHECK-NEXT: nse = 8171 // 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 = 12180 // 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 return196 }197}198