132 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=2 reassociate-fp-reductions=true enable-index-optimizations=true29// RUN: %{compile} | %{run} | FileCheck %s30//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#CSR = #sparse_tensor.encoding<{35 map = (d0, d1) -> (d0 : dense, d1 : compressed)36}>37 38#CSC = #sparse_tensor.encoding<{39 map = (d0, d1) -> (d1 : dense, d0 : compressed)40}>41 42#BSR = #sparse_tensor.encoding<{43 map = ( i, j ) ->44 ( i floordiv 2 : dense,45 j floordiv 2 : compressed,46 i mod 2 : dense,47 j mod 2 : dense48 )49}>50 51 52//53// Integration test that tests conversions between sparse tensors.54//55module {56 //57 // Main driver.58 //59 func.func @main() {60 %c0 = arith.constant 0 : index61 %c1 = arith.constant 1 : index62 %c2 = arith.constant 2 : index63 64 //65 // Initialize a 2-dim dense tensor.66 //67 %t = arith.constant sparse<[[0, 0], [0, 1], [0, 2], [0, 3],68 [1, 0], [1, 1], [1, 2], [1, 3]],69 [ 1.0, 2.0, 3.0, 4.0,70 5.0, 6.0, 7.0, 8.0 ]> : tensor<2x4xf64>71 72 %td = arith.constant dense<[[ 1.0, 2.0, 3.0, 4.0 ],73 [ 5.0, 6.0, 7.0, 8.0 ]]> : tensor<2x4xf64>74 75 // constant -> BSR (either from SparseElementAttibutes or DenseElementAttribute)76 %1 = sparse_tensor.convert %t : tensor<2x4xf64> to tensor<2x4xf64, #BSR>77 %2 = sparse_tensor.convert %td : tensor<2x4xf64> to tensor<2x4xf64, #BSR>78 %3 = sparse_tensor.convert %1 : tensor<2x4xf64, #BSR> to tensor<2x4xf64, #CSR>79 %4 = sparse_tensor.convert %1 : tensor<2x4xf64, #BSR> to tensor<2x4xf64, #CSC>80 81 //82 // CHECK: ---- Sparse Tensor ----83 // CHECK-NEXT: nse = 884 // CHECK-NEXT: dim = ( 2, 4 )85 // CHECK-NEXT: lvl = ( 1, 2, 2, 2 )86 // CHECK-NEXT: pos[1] : ( 0, 2 )87 // CHECK-NEXT: crd[1] : ( 0, 1 )88 // CHECK-NEXT: values : ( 1, 2, 5, 6, 3, 4, 7, 8 )89 // CHECK-NEXT: ----90 //91 // CHECK: ---- Sparse Tensor ----92 // CHECK-NEXT: nse = 893 // CHECK-NEXT: dim = ( 2, 4 )94 // CHECK-NEXT: lvl = ( 1, 2, 2, 2 )95 // CHECK-NEXT: pos[1] : ( 0, 2 )96 // CHECK-NEXT: crd[1] : ( 0, 1 )97 // CHECK-NEXT: values : ( 1, 2, 5, 6, 3, 4, 7, 8 )98 // CHECK-NEXT: ----99 //100 // CHECK: ---- Sparse Tensor ----101 // CHECK-NEXT: nse = 8102 // CHECK-NEXT: dim = ( 2, 4 )103 // CHECK-NEXT: lvl = ( 2, 4 )104 // CHECK-NEXT: pos[1] : ( 0, 4, 8 )105 // CHECK-NEXT: crd[1] : ( 0, 1, 2, 3, 0, 1, 2, 3 )106 // CHECK-NEXT: values : ( 1, 2, 3, 4, 5, 6, 7, 8 )107 // CHECK-NEXT: ----108 //109 // CHECK: ---- Sparse Tensor ----110 // CHECK-NEXT: nse = 8111 // CHECK-NEXT: dim = ( 2, 4 )112 // CHECK-NEXT: lvl = ( 4, 2 )113 // CHECK-NEXT: pos[1] : ( 0, 2, 4, 6, 8 )114 // CHECK-NEXT: crd[1] : ( 0, 1, 0, 1, 0, 1, 0, 1 )115 // CHECK-NEXT: values : ( 1, 5, 2, 6, 3, 7, 4, 8 )116 // CHECK-NEXT: ----117 //118 sparse_tensor.print %1 : tensor<2x4xf64, #BSR>119 sparse_tensor.print %2 : tensor<2x4xf64, #BSR>120 sparse_tensor.print %3 : tensor<2x4xf64, #CSR>121 sparse_tensor.print %4 : tensor<2x4xf64, #CSC>122 123 // TODO: Fix memory leaks.124 bufferization.dealloc_tensor %1 : tensor<2x4xf64, #BSR>125 bufferization.dealloc_tensor %2 : tensor<2x4xf64, #BSR>126 bufferization.dealloc_tensor %3 : tensor<2x4xf64, #CSR>127 bufferization.dealloc_tensor %4 : tensor<2x4xf64, #CSC>128 129 return130 }131}132