1//-------------------------------------------------------------------------------------------------- 2// WHEN CREATING A NEW TEST, PLEASE JUST COPY & PASTE WITHOUT EDITS. 3// 4// Set-up that's shared across all tests in this directory. In principle, this 5// config could be moved to lit.local.cfg. However, there are downstream users that 6// do not use these LIT config files. Hence why this is kept inline. 7// 8// DEFINE: %{sparsifier_opts} = enable-runtime-library=true 9// DEFINE: %{sparsifier_opts_sve} = enable-arm-sve=true %{sparsifier_opts} 10// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}" 11// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}" 12// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils 13// DEFINE: %{run_libs_sve} = -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils 14// DEFINE: %{run_opts} = -e main -entry-point-result=void 15// DEFINE: %{run} = mlir-runner %{run_opts} %{run_libs} 16// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs_sve} 17// 18// DEFINE: %{env} = 19//-------------------------------------------------------------------------------------------------- 20 21// RUN: %{compile} | %{run} | FileCheck %s 22// 23// Do the same run, but now with direct IR generation. 24// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false 25// RUN: %{compile} | %{run} | FileCheck %s 26// 27// Do the same run, but now with vectorization. 28// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false vl=2 reassociate-fp-reductions=true enable-index-optimizations=true 29// RUN: %{compile} | %{run} | FileCheck %s 30// 31// Do the same run, but now with VLA vectorization. 32// RUN: %if mlir_arm_sve_tests %{ %{compile_sve} | %{run_sve} | FileCheck %s %} 33 34// 35// Several common sparse storage schemes. 36// 37 38#Dense = #sparse_tensor.encoding<{ 39 map = (d0, d1) -> (d0 : dense, d1 : dense) 40}> 41 42#CSR = #sparse_tensor.encoding<{ 43 map = (d0, d1) -> (d0 : dense, d1 : compressed) 44}> 45 46#DCSR = #sparse_tensor.encoding<{ 47 map = (d0, d1) -> (d0 : compressed, d1 : compressed) 48}> 49 50#CSC = #sparse_tensor.encoding<{ 51 map = (d0, d1) -> (d1 : dense, d0 : compressed) 52}> 53 54#DCSC = #sparse_tensor.encoding<{ 55 map = (d0, d1) -> (d1 : compressed, d0 : compressed) 56}> 57 58#BlockRow = #sparse_tensor.encoding<{ 59 map = (d0, d1) -> (d0 : compressed, d1 : dense) 60}> 61 62#BlockCol = #sparse_tensor.encoding<{ 63 map = (d0, d1) -> (d1 : compressed, d0 : dense) 64}> 65 66// 67// Integration test that looks "under the hood" of sparse storage schemes. 68// 69module { 70 // 71 // Main driver that initializes a sparse tensor and inspects the sparse 72 // storage schemes in detail. Note that users of the MLIR sparsifier 73 // are typically not concerned with such details, but the test ensures 74 // everything is working "under the hood". 75 // 76 func.func @main() { 77 %c0 = arith.constant 0 : index 78 %c1 = arith.constant 1 : index 79 %d0 = arith.constant 0.0 : f64 80 81 // 82 // Initialize a dense tensor. 83 // 84 %t = arith.constant dense<[ 85 [ 1.0, 0.0, 2.0, 0.0, 0.0, 0.0, 0.0, 3.0], 86 [ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], 87 [ 0.0, 0.0, 4.0, 0.0, 0.0, 0.0, 0.0, 0.0], 88 [ 0.0, 0.0, 0.0, 5.0, 0.0, 0.0, 0.0, 0.0], 89 [ 0.0, 0.0, 0.0, 0.0, 6.0, 0.0, 0.0, 0.0], 90 [ 0.0, 7.0, 8.0, 0.0, 0.0, 0.0, 0.0, 9.0], 91 [ 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 11.0, 12.0], 92 [ 0.0, 13.0, 14.0, 0.0, 0.0, 0.0, 15.0, 16.0], 93 [ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], 94 [ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 17.0, 0.0] 95 ]> : tensor<10x8xf64> 96 97 // 98 // Convert dense tensor to various sparse tensors. 99 // 100 %0 = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #Dense> 101 %1 = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #CSR> 102 %2 = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #DCSR> 103 %3 = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #CSC> 104 %4 = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #DCSC> 105 %x = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #BlockRow> 106 %y = sparse_tensor.convert %t : tensor<10x8xf64> to tensor<10x8xf64, #BlockCol> 107 108 // 109 // Inspect storage scheme of Dense. 110 // 111 // CHECK: ---- Sparse Tensor ---- 112 // CHECK-NEXT: nse = 80 113 // CHECK-NEXT: dim = ( 10, 8 ) 114 // CHECK-NEXT: lvl = ( 10, 8 ) 115 // CHECK-NEXT: values : ( 1, 0, 2, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0, 0, 0, 7, 8, 0, 0, 0, 0, 9, 0, 0, 10, 0, 0, 0, 11, 12, 0, 13, 14, 0, 0, 0, 15, 16, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 17, 0 ) 116 // CHECK-NEXT: ---- 117 // 118 sparse_tensor.print %0 : tensor<10x8xf64, #Dense> 119 120 // 121 // Inspect storage scheme of CSR. 122 // 123 // 124 // CHECK: ---- Sparse Tensor ---- 125 // CHECK-NEXT: nse = 17 126 // CHECK-NEXT: dim = ( 10, 8 ) 127 // CHECK-NEXT: lvl = ( 10, 8 ) 128 // CHECK-NEXT: pos[1] : ( 0, 3, 3, 4, 5, 6, 9, 12, 16, 16, 17 ) 129 // CHECK-NEXT: crd[1] : ( 0, 2, 7, 2, 3, 4, 1, 2, 7, 2, 6, 7, 1, 2, 6, 7, 6 ) 130 // CHECK-NEXT: values : ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17 ) 131 // CHECK-NEXT: ---- 132 // 133 sparse_tensor.print %1 : tensor<10x8xf64, #CSR> 134 135 // 136 // Inspect storage scheme of DCSR. 137 // 138 // CHECK: ---- Sparse Tensor ---- 139 // CHECK-NEXT: nse = 17 140 // CHECK-NEXT: dim = ( 10, 8 ) 141 // CHECK-NEXT: lvl = ( 10, 8 ) 142 // CHECK-NEXT: pos[0] : ( 0, 8 ) 143 // CHECK-NEXT: crd[0] : ( 0, 2, 3, 4, 5, 6, 7, 9 ) 144 // CHECK-NEXT: pos[1] : ( 0, 3, 4, 5, 6, 9, 12, 16, 17 ) 145 // CHECK-NEXT: crd[1] : ( 0, 2, 7, 2, 3, 4, 1, 2, 7, 2, 6, 7, 1, 2, 6, 7, 6 ) 146 // CHECK-NEXT: values : ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17 ) 147 // CHECK-NEXT: ---- 148 // 149 sparse_tensor.print %2 : tensor<10x8xf64, #DCSR> 150 151 // 152 // Inspect storage scheme of CSC. 153 // 154 // CHECK: ---- Sparse Tensor ---- 155 // CHECK-NEXT: nse = 17 156 // CHECK-NEXT: dim = ( 10, 8 ) 157 // CHECK-NEXT: lvl = ( 8, 10 ) 158 // CHECK-NEXT: pos[1] : ( 0, 1, 3, 8, 9, 10, 10, 13, 17 ) 159 // CHECK-NEXT: crd[1] : ( 0, 5, 7, 0, 2, 5, 6, 7, 3, 4, 6, 7, 9, 0, 5, 6, 7 ) 160 // CHECK-NEXT: values : ( 1, 7, 13, 2, 4, 8, 10, 14, 5, 6, 11, 15, 17, 3, 9, 12, 16 ) 161 // CHECK-NEXT: ---- 162 // 163 sparse_tensor.print %3 : tensor<10x8xf64, #CSC> 164 165 // 166 // Inspect storage scheme of DCSC. 167 // 168 // CHECK: ---- Sparse Tensor ---- 169 // CHECK-NEXT: nse = 17 170 // CHECK-NEXT: dim = ( 10, 8 ) 171 // CHECK-NEXT: lvl = ( 8, 10 ) 172 // CHECK-NEXT: pos[0] : ( 0, 7 ) 173 // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3, 4, 6, 7 ) 174 // CHECK-NEXT: pos[1] : ( 0, 1, 3, 8, 9, 10, 13, 17 ) 175 // CHECK-NEXT: crd[1] : ( 0, 5, 7, 0, 2, 5, 6, 7, 3, 4, 6, 7, 9, 0, 5, 6, 7 ) 176 // CHECK-NEXT: values : ( 1, 7, 13, 2, 4, 8, 10, 14, 5, 6, 11, 15, 17, 3, 9, 12, 16 ) 177 // CHECK-NEXT: ---- 178 // 179 sparse_tensor.print %4 : tensor<10x8xf64, #DCSC> 180 181 // 182 // Inspect storage scheme of BlockRow. 183 // 184 // CHECK: ---- Sparse Tensor ---- 185 // CHECK-NEXT: nse = 64 186 // CHECK-NEXT: dim = ( 10, 8 ) 187 // CHECK-NEXT: lvl = ( 10, 8 ) 188 // CHECK-NEXT: pos[0] : ( 0, 8 ) 189 // CHECK-NEXT: crd[0] : ( 0, 2, 3, 4, 5, 6, 7, 9 ) 190 // CHECK-NEXT: values : ( 1, 0, 2, 0, 0, 0, 0, 3, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0, 0, 0, 7, 8, 0, 0, 0, 0, 9, 0, 0, 10, 0, 0, 0, 11, 12, 0, 13, 14, 0, 0, 0, 15, 16, 0, 0, 0, 0, 0, 0, 17, 0 ) 191 // CHECK-NEXT: ---- 192 // 193 sparse_tensor.print %x : tensor<10x8xf64, #BlockRow> 194 195 // 196 // Inspect storage scheme of BlockCol. 197 // 198 // CHECK: ---- Sparse Tensor ---- 199 // CHECK-NEXT: nse = 70 200 // CHECK-NEXT: dim = ( 10, 8 ) 201 // CHECK-NEXT: lvl = ( 8, 10 ) 202 // CHECK-NEXT: pos[0] : ( 0, 7 ) 203 // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3, 4, 6, 7 ) 204 // CHECK-NEXT: values : ( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7, 0, 13, 0, 0, 2, 0, 4, 0, 0, 8, 10, 14, 0, 0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 11, 15, 0, 17, 3, 0, 0, 0, 0, 9, 12, 16, 0, 0 ) 205 // CHECK-NEXT: ---- 206 // 207 sparse_tensor.print %y : tensor<10x8xf64, #BlockCol> 208 209 // Release the resources. 210 bufferization.dealloc_tensor %0 : tensor<10x8xf64, #Dense> 211 bufferization.dealloc_tensor %1 : tensor<10x8xf64, #CSR> 212 bufferization.dealloc_tensor %2 : tensor<10x8xf64, #DCSR> 213 bufferization.dealloc_tensor %3 : tensor<10x8xf64, #CSC> 214 bufferization.dealloc_tensor %4 : tensor<10x8xf64, #DCSC> 215 bufferization.dealloc_tensor %x : tensor<10x8xf64, #BlockRow> 216 bufferization.dealloc_tensor %y : tensor<10x8xf64, #BlockCol> 217 218 return 219 } 220} 221