xref: /llvm-project/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_abs.mlir (revision eb206e9ea84eff0a0596fed2de8316d924f946d1)
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 vl=2 reassociate-fp-reductions=true enable-index-optimizations=true
29// RUN: %{compile} | %{run} | FileCheck %s
30//
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#SparseVector = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
35
36#trait_op = {
37  indexing_maps = [
38    affine_map<(i) -> (i)>, // a
39    affine_map<(i) -> (i)>  // x (out)
40  ],
41  iterator_types = ["parallel"],
42  doc = "x(i) = OP a(i)"
43}
44
45module {
46  func.func @sparse_absf(%arg0: tensor<?xf64, #SparseVector>)
47                             -> tensor<?xf64, #SparseVector> {
48    %c0 = arith.constant 0 : index
49    %d = tensor.dim %arg0, %c0 : tensor<?xf64, #SparseVector>
50    %xin = tensor.empty(%d) : tensor<?xf64, #SparseVector>
51    %0 = linalg.generic #trait_op
52      ins(%arg0: tensor<?xf64, #SparseVector>)
53      outs(%xin: tensor<?xf64, #SparseVector>) {
54      ^bb0(%a: f64, %x: f64) :
55        %result = math.absf %a : f64
56        linalg.yield %result : f64
57    } -> tensor<?xf64, #SparseVector>
58    return %0 : tensor<?xf64, #SparseVector>
59  }
60
61  func.func @sparse_absi(%arg0: tensor<?xi32, #SparseVector>)
62                             -> tensor<?xi32, #SparseVector> {
63    %c0 = arith.constant 0 : index
64    %d = tensor.dim %arg0, %c0 : tensor<?xi32, #SparseVector>
65    %xin = tensor.empty(%d) : tensor<?xi32, #SparseVector>
66    %0 = linalg.generic #trait_op
67      ins(%arg0: tensor<?xi32, #SparseVector>)
68      outs(%xin: tensor<?xi32, #SparseVector>) {
69      ^bb0(%a: i32, %x: i32) :
70        %result = math.absi %a : i32
71        linalg.yield %result : i32
72    } -> tensor<?xi32, #SparseVector>
73    return %0 : tensor<?xi32, #SparseVector>
74  }
75
76  // Driver method to call and verify sign kernel.
77  func.func @main() {
78    %c0 = arith.constant 0 : index
79    %df = arith.constant 99.99 : f64
80    %di = arith.constant 9999 : i32
81
82    %pnan = arith.constant 0x7FF0000001000000 : f64
83    %nnan = arith.constant 0xFFF0000001000000 : f64
84    %pinf = arith.constant 0x7FF0000000000000 : f64
85    %ninf = arith.constant 0xFFF0000000000000 : f64
86
87    // Setup sparse vectors.
88    %v1 = arith.constant sparse<
89       [ [0], [3], [5], [11], [13], [17], [18], [20], [21], [28], [29], [31] ],
90         [ -1.5, 1.5, -10.2, 11.3, 1.0, -1.0,
91           0x7FF0000001000000, // +NaN
92           0xFFF0000001000000, // -NaN
93           0x7FF0000000000000, // +Inf
94           0xFFF0000000000000, // -Inf
95           -0.0,               // -Zero
96           0.0                 // +Zero
97        ]
98    > : tensor<32xf64>
99    %v2 = arith.constant sparse<
100       [ [0], [3], [5], [11], [13], [17], [18], [21], [31] ],
101         [ -2147483648, -2147483647, -1000, -1, 0,
102           1, 1000, 2147483646, 2147483647
103         ]
104    > : tensor<32xi32>
105    %sv1 = sparse_tensor.convert %v1
106         : tensor<32xf64> to tensor<?xf64, #SparseVector>
107    %sv2 = sparse_tensor.convert %v2
108         : tensor<32xi32> to tensor<?xi32, #SparseVector>
109
110    // Call abs kernels.
111    %0 = call @sparse_absf(%sv1) : (tensor<?xf64, #SparseVector>)
112                                 -> tensor<?xf64, #SparseVector>
113
114    %1 = call @sparse_absi(%sv2) : (tensor<?xi32, #SparseVector>)
115                                 -> tensor<?xi32, #SparseVector>
116
117    //
118    // Verify the results.
119    //
120    // CHECK:      ---- Sparse Tensor ----
121    // CHECK-NEXT: nse = 12
122    // CHECK-NEXT: dim = ( 32 )
123    // CHECK-NEXT: lvl = ( 32 )
124    // CHECK-NEXT: pos[0] : ( 0, 12 )
125    // CHECK-NEXT: crd[0] : ( 0, 3, 5, 11, 13, 17, 18, 20, 21, 28, 29, 31 )
126    // CHECK-NEXT: values : ( 1.5, 1.5, 10.2, 11.3, 1, 1, nan, nan, inf, inf, 0, 0 )
127    // CHECK-NEXT: ----
128    //
129    // CHECK-NEXT: ---- Sparse Tensor ----
130    // CHECK-NEXT: nse = 9
131    // CHECK-NEXT: dim = ( 32 )
132    // CHECK-NEXT: lvl = ( 32 )
133    // CHECK-NEXT: pos[0] : ( 0, 9 )
134    // CHECK-NEXT: crd[0] : ( 0, 3, 5, 11, 13, 17, 18, 21, 31 )
135    // CHECK-NEXT: values : ( -2147483648, 2147483647, 1000, 1, 0, 1, 1000, 2147483646, 2147483647 )
136    // CHECK-NEXT: ----
137    //
138    sparse_tensor.print %0 : tensor<?xf64, #SparseVector>
139    sparse_tensor.print %1 : tensor<?xi32, #SparseVector>
140
141    // Release the resources.
142    bufferization.dealloc_tensor %sv1 : tensor<?xf64, #SparseVector>
143    bufferization.dealloc_tensor %sv2 : tensor<?xi32, #SparseVector>
144    bufferization.dealloc_tensor %0 : tensor<?xf64, #SparseVector>
145    bufferization.dealloc_tensor %1 : tensor<?xi32, #SparseVector>
146    return
147  }
148}
149