175 lines · plain
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, this5// config could be moved to lit.local.cfg. However, there are downstream users that6// do not use these LIT config files. Hence why this is kept inline.7//8// DEFINE: %{sparsifier_opts} = enable-runtime-library=true9// 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_utils13// DEFINE: %{run_libs_sve} = -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils14// DEFINE: %{run_opts} = -e main -entry-point-result=void15// 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 %s22//23// Do the same run, but now with direct IR generation.24// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false enable-buffer-initialization=true25// RUN: %{compile} | %{run} | FileCheck %s26//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=4 reassociate-fp-reductions=true enable-index-optimizations=true29// RUN: %{compile} | %{run} | FileCheck %s30 31#MAT_C_C = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : compressed, d1 : compressed)}>32#MAT_D_C = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : dense, d1 : compressed)}>33#MAT_C_D = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : compressed, d1 : dense)}>34#MAT_D_D = #sparse_tensor.encoding<{35 map = (d0, d1) -> (d1 : dense, d0 : dense)36}>37 38#MAT_C_C_P = #sparse_tensor.encoding<{39 map = (d0, d1) -> (d1 : compressed, d0 : compressed)40}>41 42#MAT_C_D_P = #sparse_tensor.encoding<{43 map = (d0, d1) -> (d1 : compressed, d0 : dense)44}>45 46#MAT_D_C_P = #sparse_tensor.encoding<{47 map = (d0, d1) -> (d1 : dense, d0 : compressed)48}>49 50module {51 func.func private @printMemrefF64(%ptr : tensor<*xf64>)52 func.func private @printMemref1dF64(%ptr : memref<?xf64>) attributes { llvm.emit_c_interface }53 54 //55 // Tests without permutation (concatenate on dimension 1)56 //57 58 // Concats all sparse matrices (with different encodings) to a sparse matrix.59 func.func @concat_sparse_sparse_dim1(%arg0: tensor<4x2xf64, #MAT_C_C>, %arg1: tensor<4x3xf64, #MAT_C_D>, %arg2: tensor<4x4xf64, #MAT_D_C>) -> tensor<4x9xf64, #MAT_C_C> {60 %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 1 : index}61 : tensor<4x2xf64, #MAT_C_C>, tensor<4x3xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C> to tensor<4x9xf64, #MAT_C_C>62 return %0 : tensor<4x9xf64, #MAT_C_C>63 }64 65 // Concats all sparse matrices (with different encodings) to a dense matrix.66 func.func @concat_sparse_dense_dim1(%arg0: tensor<4x2xf64, #MAT_C_C>, %arg1: tensor<4x3xf64, #MAT_C_D>, %arg2: tensor<4x4xf64, #MAT_D_C>) -> tensor<4x9xf64> {67 %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 1 : index}68 : tensor<4x2xf64, #MAT_C_C>, tensor<4x3xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C> to tensor<4x9xf64>69 return %0 : tensor<4x9xf64>70 }71 72 // Concats mix sparse and dense matrices to a sparse matrix.73 func.func @concat_mix_sparse_dim1(%arg0: tensor<4x2xf64>, %arg1: tensor<4x3xf64, #MAT_C_D>, %arg2: tensor<4x4xf64, #MAT_D_C>) -> tensor<4x9xf64, #MAT_C_C> {74 %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 1 : index}75 : tensor<4x2xf64>, tensor<4x3xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C> to tensor<4x9xf64, #MAT_C_C>76 return %0 : tensor<4x9xf64, #MAT_C_C>77 }78 79 // Concats mix sparse and dense matrices to a dense matrix.80 func.func @concat_mix_dense_dim1(%arg0: tensor<4x2xf64>, %arg1: tensor<4x3xf64, #MAT_C_D>, %arg2: tensor<4x4xf64, #MAT_D_C>) -> tensor<4x9xf64> {81 %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 1 : index}82 : tensor<4x2xf64>, tensor<4x3xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C> to tensor<4x9xf64>83 return %0 : tensor<4x9xf64>84 }85 86 func.func @dump_mat_dense_4x9(%A: tensor<4x9xf64>) {87 %1 = tensor.cast %A : tensor<4x9xf64> to tensor<*xf64>88 call @printMemrefF64(%1) : (tensor<*xf64>) -> ()89 90 return91 }92 93 // Driver method to call and verify kernels.94 func.func @main() {95 %m42 = arith.constant dense<96 [ [ 1.0, 0.0 ],97 [ 3.1, 0.0 ],98 [ 0.0, 2.0 ],99 [ 0.0, 0.0 ] ]> : tensor<4x2xf64>100 %m43 = arith.constant dense<101 [ [ 1.0, 0.0, 1.0 ],102 [ 1.0, 0.0, 0.5 ],103 [ 0.0, 0.0, 1.0 ],104 [ 5.0, 2.0, 0.0 ] ]> : tensor<4x3xf64>105 %m44 = arith.constant dense<106 [ [ 0.0, 0.0, 1.5, 1.0],107 [ 0.0, 3.5, 0.0, 0.0],108 [ 1.0, 5.0, 2.0, 0.0],109 [ 1.0, 0.5, 0.0, 0.0] ]> : tensor<4x4xf64>110 111 %sm42cc = sparse_tensor.convert %m42 : tensor<4x2xf64> to tensor<4x2xf64, #MAT_C_C>112 %sm43cd = sparse_tensor.convert %m43 : tensor<4x3xf64> to tensor<4x3xf64, #MAT_C_D>113 %sm44dc = sparse_tensor.convert %m44 : tensor<4x4xf64> to tensor<4x4xf64, #MAT_D_C>114 115 //116 // CHECK: ---- Sparse Tensor ----117 // CHECK-NEXT: nse = 18118 // CHECK-NEXT: dim = ( 4, 9 )119 // CHECK-NEXT: lvl = ( 4, 9 )120 // CHECK-NEXT: pos[0] : ( 0, 4 )121 // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3 )122 // CHECK-NEXT: pos[1] : ( 0, 5, 9, 14, 18 )123 // CHECK-NEXT: crd[1] : ( 0, 2, 4, 7, 8, 0, 2, 4, 6, 1, 4, 5, 6, 7, 2, 3, 5, 6 )124 // CHECK-NEXT: values : ( 1, 1, 1, 1.5, 1, 3.1, 1, 0.5, 3.5, 2, 1, 1, 5, 2, 5, 2, 1, 0.5 )125 // CHECK-NEXT: ----126 //127 %8 = call @concat_sparse_sparse_dim1(%sm42cc, %sm43cd, %sm44dc)128 : (tensor<4x2xf64, #MAT_C_C>, tensor<4x3xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C>) -> tensor<4x9xf64, #MAT_C_C>129 sparse_tensor.print %8 : tensor<4x9xf64, #MAT_C_C>130 131 // CHECK: {{\[}}[1, 0, 1, 0, 1, 0, 0, 1.5, 1],132 // CHECK-NEXT: [3.1, 0, 1, 0, 0.5, 0, 3.5, 0, 0],133 // CHECK-NEXT: [0, 2, 0, 0, 1, 1, 5, 2, 0],134 // CHECK-NEXT: [0, 0, 5, 2, 0, 1, 0.5, 0, 0]]135 %9 = call @concat_sparse_dense_dim1(%sm42cc, %sm43cd, %sm44dc)136 : (tensor<4x2xf64, #MAT_C_C>, tensor<4x3xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C>) -> tensor<4x9xf64>137 call @dump_mat_dense_4x9(%9) : (tensor<4x9xf64>) -> ()138 139 //140 // CHECK: ---- Sparse Tensor ----141 // CHECK-NEXT: nse = 18142 // CHECK-NEXT: dim = ( 4, 9 )143 // CHECK-NEXT: lvl = ( 4, 9 )144 // CHECK-NEXT: pos[0] : ( 0, 4 )145 // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3 )146 // CHECK-NEXT: pos[1] : ( 0, 5, 9, 14, 18 )147 // CHECK-NEXT: crd[1] : ( 0, 2, 4, 7, 8, 0, 2, 4, 6, 1, 4, 5, 6, 7, 2, 3, 5, 6 )148 // CHECK-NEXT: values : ( 1, 1, 1, 1.5, 1, 3.1, 1, 0.5, 3.5, 2, 1, 1, 5, 2, 5, 2, 1, 0.5 )149 // CHECK-NEXT: ----150 //151 %10 = call @concat_mix_sparse_dim1(%m42, %sm43cd, %sm44dc)152 : (tensor<4x2xf64>, tensor<4x3xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C>) -> tensor<4x9xf64, #MAT_C_C>153 sparse_tensor.print %10 : tensor<4x9xf64, #MAT_C_C>154 155 // CHECK: {{\[}}[1, 0, 1, 0, 1, 0, 0, 1.5, 1],156 // CHECK-NEXT: [3.1, 0, 1, 0, 0.5, 0, 3.5, 0, 0],157 // CHECK-NEXT: [0, 2, 0, 0, 1, 1, 5, 2, 0],158 // CHECK-NEXT: [0, 0, 5, 2, 0, 1, 0.5, 0, 0]]159 %11 = call @concat_mix_dense_dim1(%m42, %sm43cd, %sm44dc)160 : (tensor<4x2xf64>, tensor<4x3xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C>) -> tensor<4x9xf64>161 call @dump_mat_dense_4x9(%11) : (tensor<4x9xf64>) -> ()162 163 // Release resources.164 bufferization.dealloc_tensor %sm42cc : tensor<4x2xf64, #MAT_C_C>165 bufferization.dealloc_tensor %sm43cd : tensor<4x3xf64, #MAT_C_D>166 bufferization.dealloc_tensor %sm44dc : tensor<4x4xf64, #MAT_D_C>167 168 bufferization.dealloc_tensor %8 : tensor<4x9xf64, #MAT_C_C>169 bufferization.dealloc_tensor %9 : tensor<4x9xf64>170 bufferization.dealloc_tensor %10 : tensor<4x9xf64, #MAT_C_C>171 bufferization.dealloc_tensor %11 : tensor<4x9xf64>172 return173 }174}175