138 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// Test that test-bufferization-analysis-only works. This option is useful35// for understanding why buffer copies were inserted.36// RUN: mlir-opt %s --sparsifier="test-bufferization-analysis-only" -o /dev/null37 38#Sparse1 = #sparse_tensor.encoding<{39 map = (i, j, k) -> (40 j : compressed,41 k : compressed,42 i : dense43 )44}>45 46#Sparse2 = #sparse_tensor.encoding<{47 map = (i, j, k) -> (48 i floordiv 2 : compressed,49 j floordiv 2 : compressed,50 k floordiv 2 : compressed,51 i mod 2 : dense,52 j mod 2 : dense,53 k mod 2 : dense)54}>55 56module {57 58 //59 // Main driver that tests sparse tensor storage.60 //61 func.func @main() {62 %c0 = arith.constant 0 : index63 %i0 = arith.constant 0 : i3264 65 // Setup input dense tensor and convert to two sparse tensors.66 %d = arith.constant dense <[67 [ // i=068 [ 1, 0, 0, 0 ],69 [ 0, 0, 0, 0 ],70 [ 0, 0, 0, 0 ],71 [ 0, 0, 5, 0 ] ],72 [ // i=173 [ 2, 0, 0, 0 ],74 [ 0, 0, 0, 0 ],75 [ 0, 0, 0, 0 ],76 [ 0, 0, 6, 0 ] ],77 [ //i=278 [ 3, 0, 0, 0 ],79 [ 0, 0, 0, 0 ],80 [ 0, 0, 0, 0 ],81 [ 0, 0, 7, 0 ] ],82 //i=383 [ [ 4, 0, 0, 0 ],84 [ 0, 0, 0, 0 ],85 [ 0, 0, 0, 0 ],86 [ 0, 0, 8, 0 ] ]87 ]> : tensor<4x4x4xi32>88 89 %a = sparse_tensor.convert %d : tensor<4x4x4xi32> to tensor<4x4x4xi32, #Sparse1>90 %b = sparse_tensor.convert %d : tensor<4x4x4xi32> to tensor<4x4x4xi32, #Sparse2>91 92 //93 // If we store the two "fibers" [1,2,3,4] starting at index (0,0,0) and94 // ending at index (3,0,0) and [5,6,7,8] starting at index (0,3,2) and95 // ending at index (3,3,2)) with a “DCSR-flavored” along (j,k) with96 // dense “fibers” in the i-dim, we end up with 8 stored entries.97 //98 // CHECK: ---- Sparse Tensor ----99 // CHECK-NEXT: nse = 8100 // CHECK-NEXT: dim = ( 4, 4, 4 )101 // CHECK-NEXT: lvl = ( 4, 4, 4 )102 // CHECK-NEXT: pos[0] : ( 0, 2 )103 // CHECK-NEXT: crd[0] : ( 0, 3 )104 // CHECK-NEXT: pos[1] : ( 0, 1, 2 )105 // CHECK-NEXT: crd[1] : ( 0, 2 )106 // CHECK-NEXT: values : ( 1, 2, 3, 4, 5, 6, 7, 8 )107 // CHECK-NEXT: ----108 //109 sparse_tensor.print %a : tensor<4x4x4xi32, #Sparse1>110 111 //112 // If we store full 2x2x2 3-D blocks in the original index order113 // in a compressed fashion, we end up with 4 blocks to incorporate114 // all the nonzeros, and thus 32 stored entries.115 //116 // CHECK: ---- Sparse Tensor ----117 // CHECK-NEXT: nse = 32118 // CHECK-NEXT: dim = ( 4, 4, 4 )119 // CHECK-NEXT: lvl = ( 2, 2, 2, 2, 2, 2 )120 // CHECK-NEXT: pos[0] : ( 0, 2 )121 // CHECK-NEXT: crd[0] : ( 0, 1 )122 // CHECK-NEXT: pos[1] : ( 0, 2, 4 )123 // CHECK-NEXT: crd[1] : ( 0, 1, 0, 1 )124 // CHECK-NEXT: pos[2] : ( 0, 1, 2, 3, 4 )125 // CHECK-NEXT: crd[2] : ( 0, 1, 0, 1 )126 // CHECK-NEXT: values : ( 1, 0, 0, 0, 2, 0, 0, 0, 0, 0, 5, 0, 0, 0, 6, 0, 3, 0, 0, 0, 4, 0, 0, 0, 0, 0, 7, 0, 0, 0, 8, 0 )127 // CHECK-NEXT: ----128 //129 sparse_tensor.print %b : tensor<4x4x4xi32, #Sparse2>130 131 // Release the resources.132 bufferization.dealloc_tensor %a : tensor<4x4x4xi32, #Sparse1>133 bufferization.dealloc_tensor %b : tensor<4x4x4xi32, #Sparse2>134 135 return136 }137}138