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=4 reassociate-fp-reductions=true enable-index-optimizations=true 29// RUN: %{compile} | %{run} | FileCheck %s 30 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 return 91 } 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 = 18 118 // 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 = 18 142 // 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 return 173 } 174} 175