126 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 do sparsification using sparse-iterator-based loops.24// REDEFINE: %{sparsifier_opts} = sparse-emit-strategy=sparse-iterator25// RUN: %{compile} | %{run} | FileCheck %s26//27 28#COO = #sparse_tensor.encoding<{29 map = (d0, d1, d2, d3) -> (30 d0 : compressed(nonunique),31 d1 : singleton(nonunique, soa),32 d2 : singleton(nonunique, soa),33 d3 : singleton(soa)34 ),35 explicitVal = 1 : i3236}>37 38#VEC = #sparse_tensor.encoding<{39 map = (d0) -> (d0 : compressed)40}>41 42 43module {44 // An example of vector reductions (lowered through sparse_tensor.iterate).45 func.func @sqsum(%arg0: tensor<2x3x4x5xi32, #COO>) -> tensor<i32> {46 %cst = arith.constant dense<0> : tensor<i32>47 %0 = linalg.generic {48 indexing_maps = [49 affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,50 affine_map<(d0, d1, d2, d3) -> ()>51 ],52 iterator_types = ["reduction", "reduction", "reduction", "reduction"]53 } ins(%arg0 : tensor<2x3x4x5xi32, #COO>) outs(%cst : tensor<i32>) {54 ^bb0(%in: i32, %out: i32):55 %1 = arith.muli %in, %in : i3256 %2 = arith.addi %out, %1 : i3257 linalg.yield %2 : i3258 } -> tensor<i32>59 return %0 : tensor<i32>60 }61 62 // An example of vector addition (lowered through sparse_tensor.coiterate).63 func.func @vec_add(%arg0: tensor<4xi32, #VEC>, %arg1: tensor<4xi32, #VEC>) -> tensor<4xi32> {64 %cst = arith.constant dense<0> : tensor<4xi32>65 %0 = linalg.generic {66 indexing_maps = [67 affine_map<(d0) -> (d0)>,68 affine_map<(d0) -> (d0)>,69 affine_map<(d0) -> (d0)>70 ],71 iterator_types = ["parallel"]72 }73 ins(%arg0, %arg1 : tensor<4xi32, #VEC>, tensor<4xi32, #VEC>)74 outs(%cst : tensor<4xi32>) {75 ^bb0(%in1: i32, %in2: i32, %out: i32):76 %2 = arith.addi %in1, %in2 : i3277 linalg.yield %2 : i3278 } -> tensor<4xi32>79 return %0 : tensor<4xi32>80 }81 82 func.func @main() {83 %c0 = arith.constant 0 : index84 %i0 = arith.constant 0 : i3285 86 %cst = arith.constant sparse<87 [88 [0, 1, 2, 3],89 [1, 1, 2, 3],90 [1, 2, 2, 3],91 [1, 2, 3, 4]92 ],93 [1, 1, 1, 1]94 > : tensor<2x3x4x5xi32>95 96 %l = arith.constant dense<97 [0, 1, 2, 3]98 > : tensor<4xi32>99 %r = arith.constant dense<100 [1, 0, 3, 0]101 > : tensor<4xi32>102 103 %input = sparse_tensor.convert %cst : tensor<2x3x4x5xi32> to tensor<2x3x4x5xi32, #COO>104 %0 = call @sqsum(%input) : (tensor<2x3x4x5xi32, #COO>) -> tensor<i32>105 %v = tensor.extract %0[] : tensor<i32>106 107 %lhs = sparse_tensor.convert %l : tensor<4xi32> to tensor<4xi32, #VEC>108 %rhs = sparse_tensor.convert %r : tensor<4xi32> to tensor<4xi32, #VEC>109 %add = call @vec_add(%lhs, %rhs) : (tensor<4xi32, #VEC>, tensor<4xi32, #VEC>) -> tensor<4xi32>110 111 // CHECK: 4112 vector.print %v : i32113 // CHECK-NEXT: ( 1, 1, 5, 3 )114 %vec = vector.transfer_read %add[%c0], %i0 : tensor<4xi32>, vector<4xi32>115 vector.print %vec : vector<4xi32>116 117 bufferization.dealloc_tensor %input : tensor<2x3x4x5xi32, #COO>118 bufferization.dealloc_tensor %0 : tensor<i32>119 120 bufferization.dealloc_tensor %lhs : tensor<4xi32, #VEC>121 bufferization.dealloc_tensor %rhs : tensor<4xi32, #VEC>122 bufferization.dealloc_tensor %add : tensor<4xi32>123 return124 }125}126