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// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false 22 23// RUN: %{compile} | %{run} | FileCheck %s 24// 25// Do the same run, but now with vectorization. 26// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false vl=2 reassociate-fp-reductions=true enable-index-optimizations=true 27// RUN: %{compile} | %{run} | FileCheck %s 28// 29// Do the same run, but now with VLA vectorization. 30// RUN: %if mlir_arm_sve_tests %{ %{compile_sve} | %{run_sve} | FileCheck %s %} 31 32#Tensor1 = #sparse_tensor.encoding<{ 33 map = (d0, d1, d2) -> (d0 : compressed(nonunique), d1 : singleton(nonunique), d2 : singleton) 34}> 35 36#Tensor2 = #sparse_tensor.encoding<{ 37 map = (d0, d1, d2) -> (d0 : dense, d1 : compressed, d2 : dense) 38}> 39 40#Tensor3 = #sparse_tensor.encoding<{ 41 map = (d0, d1, d2) -> (d0 : dense, d2 : dense, d1 : compressed) 42}> 43 44module { 45 // 46 // Utility for output. 47 // 48 func.func @dump(%arg0: tensor<2x3x4xf32>) { 49 %c0 = arith.constant 0 : index 50 %d0 = arith.constant -1.0 : f32 51 %0 = vector.transfer_read %arg0[%c0, %c0, %c0], %d0: tensor<2x3x4xf32>, vector<2x3x4xf32> 52 vector.print %0 : vector<2x3x4xf32> 53 return 54 } 55 56 // 57 // The first test suite (for non-singleton LevelTypes). 58 // 59 func.func @main() { 60 // 61 // Initialize a 3-dim dense tensor. 62 // 63 %src = arith.constant dense<[ 64 [ [ 1.0, 2.0, 3.0, 4.0 ], 65 [ 5.0, 6.0, 7.0, 8.0 ], 66 [ 9.0, 10.0, 11.0, 12.0 ] ], 67 [ [ 13.0, 14.0, 15.0, 16.0 ], 68 [ 17.0, 18.0, 19.0, 20.0 ], 69 [ 21.0, 22.0, 23.0, 24.0 ] ] 70 ]> : tensor<2x3x4xf64> 71 72 // 73 // Convert dense tensor directly to various sparse tensors. 74 // 75 %s1 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor1> 76 %s2 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor2> 77 %s3 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor3> 78 79 // 80 // Convert sparse tensor directly to another sparse format. 81 // 82 %t1 = sparse_tensor.convert %s1 : tensor<2x3x4xf64, #Tensor1> to tensor<2x3x4xf32, #Tensor1> 83 %t2 = sparse_tensor.convert %s2 : tensor<2x3x4xf64, #Tensor2> to tensor<2x3x4xf32, #Tensor2> 84 %t3 = sparse_tensor.convert %s3 : tensor<2x3x4xf64, #Tensor3> to tensor<2x3x4xf32, #Tensor3> 85 86 // 87 // Convert sparse tensor back to dense. 88 // 89 %d1 = sparse_tensor.convert %t1 : tensor<2x3x4xf32, #Tensor1> to tensor<2x3x4xf32> 90 %d2 = sparse_tensor.convert %t2 : tensor<2x3x4xf32, #Tensor2> to tensor<2x3x4xf32> 91 %d3 = sparse_tensor.convert %t3 : tensor<2x3x4xf32, #Tensor3> to tensor<2x3x4xf32> 92 93 // 94 // Check round-trip equality. And release dense tensors. 95 // 96 // CHECK-COUNT-3: ( ( ( 1, 2, 3, 4 ), ( 5, 6, 7, 8 ), ( 9, 10, 11, 12 ) ), ( ( 13, 14, 15, 16 ), ( 17, 18, 19, 20 ), ( 21, 22, 23, 24 ) ) ) 97 call @dump(%d1) : (tensor<2x3x4xf32>) -> () 98 call @dump(%d2) : (tensor<2x3x4xf32>) -> () 99 call @dump(%d3) : (tensor<2x3x4xf32>) -> () 100 101 // 102 // Release sparse tensors. 103 // 104 bufferization.dealloc_tensor %t1 : tensor<2x3x4xf32, #Tensor1> 105 bufferization.dealloc_tensor %t2 : tensor<2x3x4xf32, #Tensor2> 106 bufferization.dealloc_tensor %t3 : tensor<2x3x4xf32, #Tensor3> 107 bufferization.dealloc_tensor %s1 : tensor<2x3x4xf64, #Tensor1> 108 bufferization.dealloc_tensor %s2 : tensor<2x3x4xf64, #Tensor2> 109 bufferization.dealloc_tensor %s3 : tensor<2x3x4xf64, #Tensor3> 110 bufferization.dealloc_tensor %d1 : tensor<2x3x4xf32> 111 bufferization.dealloc_tensor %d2 : tensor<2x3x4xf32> 112 bufferization.dealloc_tensor %d3 : tensor<2x3x4xf32> 113 114 return 115 } 116} 117