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 enable-buffer-initialization=true 25// RUN: %{compile} | %{run} | FileCheck %s 26// 27// Do the same run, but now with direct IR generation and vectorization. 28// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false enable-buffer-initialization=true 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 direct IR generation and VLA vectorization. 32// RUN: %if mlir_arm_sve_tests %{ %{compile_sve} | %{run_sve} | FileCheck %s %} 33 34#MAT_C_C = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : compressed, d1 : compressed)}> 35#MAT_D_C = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : dense, d1 : compressed)}> 36#MAT_C_D = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : compressed, d1 : dense)}> 37#MAT_D_D = #sparse_tensor.encoding<{ 38 map = (d0, d1) -> (d1 : dense, d0 : dense) 39}> 40 41#MAT_C_C_P = #sparse_tensor.encoding<{ 42 map = (d0, d1) -> (d1 : compressed, d0 : compressed) 43}> 44 45#MAT_C_D_P = #sparse_tensor.encoding<{ 46 map = (d0, d1) -> (d1 : compressed, d0 : dense), 47}> 48 49#MAT_D_C_P = #sparse_tensor.encoding<{ 50 map = (d0, d1) -> (d1 : dense, d0 : compressed) 51}> 52 53module { 54 func.func private @printMemrefF64(%ptr : tensor<*xf64>) 55 56 // Concats all sparse matrices (with different encodings) to a sparse matrix. 57 func.func @concat_sparse_sparse(%arg0: tensor<2x4xf64, #MAT_C_C>, %arg1: tensor<3x4xf64, #MAT_C_D>, %arg2: tensor<4x4xf64, #MAT_D_C>) -> tensor<9x4xf64, #MAT_C_C> { 58 %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index} 59 : tensor<2x4xf64, #MAT_C_C>, tensor<3x4xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C> to tensor<9x4xf64, #MAT_C_C> 60 return %0 : tensor<9x4xf64, #MAT_C_C> 61 } 62 63 // Concats all sparse matrices (with different encodings) to a dense matrix. 64 func.func @concat_sparse_dense(%arg0: tensor<2x4xf64, #MAT_C_C>, %arg1: tensor<3x4xf64, #MAT_C_D>, %arg2: tensor<4x4xf64, #MAT_D_C>) -> tensor<9x4xf64> { 65 %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index} 66 : tensor<2x4xf64, #MAT_C_C>, tensor<3x4xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C> to tensor<9x4xf64> 67 return %0 : tensor<9x4xf64> 68 } 69 70 // Concats mix sparse and dense matrices to a sparse matrix. 71 func.func @concat_mix_sparse(%arg0: tensor<2x4xf64>, %arg1: tensor<3x4xf64, #MAT_C_D>, %arg2: tensor<4x4xf64, #MAT_D_C>) -> tensor<9x4xf64, #MAT_C_C> { 72 %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index} 73 : tensor<2x4xf64>, tensor<3x4xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C> to tensor<9x4xf64, #MAT_C_C> 74 return %0 : tensor<9x4xf64, #MAT_C_C> 75 } 76 77 // Concats mix sparse and dense matrices to a dense matrix. 78 func.func @concat_mix_dense(%arg0: tensor<2x4xf64>, %arg1: tensor<3x4xf64, #MAT_C_D>, %arg2: tensor<4x4xf64, #MAT_D_C>) -> tensor<9x4xf64> { 79 %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index} 80 : tensor<2x4xf64>, tensor<3x4xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C> to tensor<9x4xf64> 81 return %0 : tensor<9x4xf64> 82 } 83 84 // Outputs dense matrix. 85 func.func @dump_mat_dense_9x4(%A: tensor<9x4xf64>) { 86 %u = tensor.cast %A : tensor<9x4xf64> to tensor<*xf64> 87 call @printMemrefF64(%u) : (tensor<*xf64>) -> () 88 return 89 } 90 91 // Driver method to call and verify kernels. 92 func.func @main() { 93 %m24 = arith.constant dense< 94 [ [ 1.0, 0.0, 3.0, 0.0], 95 [ 0.0, 2.0, 0.0, 0.0] ]> : tensor<2x4xf64> 96 %m34 = arith.constant dense< 97 [ [ 1.0, 0.0, 1.0, 1.0], 98 [ 0.0, 0.5, 0.0, 0.0], 99 [ 1.0, 5.0, 2.0, 0.0] ]> : tensor<3x4xf64> 100 %m44 = arith.constant dense< 101 [ [ 0.0, 0.0, 1.5, 1.0], 102 [ 0.0, 3.5, 0.0, 0.0], 103 [ 1.0, 5.0, 2.0, 0.0], 104 [ 1.0, 0.5, 0.0, 0.0] ]> : tensor<4x4xf64> 105 106 %sm24cc = sparse_tensor.convert %m24 : tensor<2x4xf64> to tensor<2x4xf64, #MAT_C_C> 107 %sm34cd = sparse_tensor.convert %m34 : tensor<3x4xf64> to tensor<3x4xf64, #MAT_C_D> 108 %sm44dc = sparse_tensor.convert %m44 : tensor<4x4xf64> to tensor<4x4xf64, #MAT_D_C> 109 110 // 111 // CHECK: ---- Sparse Tensor ---- 112 // CHECK-NEXT: nse = 18 113 // CHECK-NEXT: dim = ( 9, 4 ) 114 // CHECK-NEXT: lvl = ( 9, 4 ) 115 // CHECK-NEXT: pos[0] : ( 0, 9 ) 116 // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3, 4, 5, 6, 7, 8 ) 117 // CHECK-NEXT: pos[1] : ( 0, 2, 3, 6, 7, 10, 12, 13, 16, 18 ) 118 // CHECK-NEXT: crd[1] : ( 0, 2, 1, 0, 2, 3, 1, 0, 1, 2, 2, 3, 1, 0, 1, 2, 0, 1 ) 119 // CHECK-NEXT: values : ( 1, 3, 2, 1, 1, 1, 0.5, 1, 5, 2, 1.5, 1, 3.5, 1, 5, 2, 1, 0.5 ) 120 // CHECK-NEXT: ---- 121 // 122 %0 = call @concat_sparse_sparse(%sm24cc, %sm34cd, %sm44dc) 123 : (tensor<2x4xf64, #MAT_C_C>, tensor<3x4xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C>) -> tensor<9x4xf64, #MAT_C_C> 124 sparse_tensor.print %0 : tensor<9x4xf64, #MAT_C_C> 125 126 // 127 // CHECK: {{\[}}[1, 0, 3, 0], 128 // CHECK-NEXT: [0, 2, 0, 0], 129 // CHECK-NEXT: [1, 0, 1, 1], 130 // CHECK-NEXT: [0, 0.5, 0, 0], 131 // CHECK-NEXT: [1, 5, 2, 0], 132 // CHECK-NEXT: [0, 0, 1.5, 1], 133 // CHECK-NEXT: [0, 3.5, 0, 0], 134 // CHECK-NEXT: [1, 5, 2, 0], 135 // CHECK-NEXT: [1, 0.5, 0, 0]] 136 // 137 %1 = call @concat_sparse_dense(%sm24cc, %sm34cd, %sm44dc) 138 : (tensor<2x4xf64, #MAT_C_C>, tensor<3x4xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C>) -> tensor<9x4xf64> 139 call @dump_mat_dense_9x4(%1) : (tensor<9x4xf64>) -> () 140 141 // 142 // CHECK: ---- Sparse Tensor ---- 143 // CHECK-NEXT: nse = 18 144 // CHECK-NEXT: dim = ( 9, 4 ) 145 // CHECK-NEXT: lvl = ( 9, 4 ) 146 // CHECK-NEXT: pos[0] : ( 0, 9 ) 147 // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3, 4, 5, 6, 7, 8 ) 148 // CHECK-NEXT: pos[1] : ( 0, 2, 3, 6, 7, 10, 12, 13, 16, 18 ) 149 // CHECK-NEXT: crd[1] : ( 0, 2, 1, 0, 2, 3, 1, 0, 1, 2, 2, 3, 1, 0, 1, 2, 0, 1 ) 150 // CHECK-NEXT: values : ( 1, 3, 2, 1, 1, 1, 0.5, 1, 5, 2, 1.5, 1, 3.5, 1, 5, 2, 1, 0.5 ) 151 // CHECK-NEXT: ---- 152 // 153 %2 = call @concat_mix_sparse(%m24, %sm34cd, %sm44dc) 154 : (tensor<2x4xf64>, tensor<3x4xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C>) -> tensor<9x4xf64, #MAT_C_C> 155 sparse_tensor.print %2 : tensor<9x4xf64, #MAT_C_C> 156 157 // 158 // CHECK: {{\[}}[1, 0, 3, 0], 159 // CHECK-NEXT: [0, 2, 0, 0], 160 // CHECK-NEXT: [1, 0, 1, 1], 161 // CHECK-NEXT: [0, 0.5, 0, 0], 162 // CHECK-NEXT: [1, 5, 2, 0], 163 // CHECK-NEXT: [0, 0, 1.5, 1], 164 // CHECK-NEXT: [0, 3.5, 0, 0], 165 // CHECK-NEXT: [1, 5, 2, 0], 166 // CHECK-NEXT: [1, 0.5, 0, 0]] 167 // 168 %3 = call @concat_mix_dense(%m24, %sm34cd, %sm44dc) 169 : (tensor<2x4xf64>, tensor<3x4xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C>) -> tensor<9x4xf64> 170 call @dump_mat_dense_9x4(%3) : (tensor<9x4xf64>) -> () 171 172 // Release resources. 173 bufferization.dealloc_tensor %sm24cc : tensor<2x4xf64, #MAT_C_C> 174 bufferization.dealloc_tensor %sm34cd : tensor<3x4xf64, #MAT_C_D> 175 bufferization.dealloc_tensor %sm44dc : tensor<4x4xf64, #MAT_D_C> 176 bufferization.dealloc_tensor %0 : tensor<9x4xf64, #MAT_C_C> 177 bufferization.dealloc_tensor %1 : tensor<9x4xf64> 178 bufferization.dealloc_tensor %2 : tensor<9x4xf64, #MAT_C_C> 179 bufferization.dealloc_tensor %3 : tensor<9x4xf64> 180 return 181 } 182} 183