1// RUN: mlir-opt %s --sparse-compiler | \
2// RUN: mlir-cpu-runner \
3// RUN:  -e entry -entry-point-result=void  \
4// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
5// RUN: FileCheck %s
6//
7// Do the same run, but now with SIMDization as well. This should not change the outcome.
8//
9// RUN: mlir-opt %s \
10// RUN:   --linalg-generalize-named-ops --linalg-fuse-elementwise-ops \
11// RUN:   --sparse-compiler="vectorization-strategy=2 vl=8" | \
12// RUN: mlir-cpu-runner \
13// RUN:  -e entry -entry-point-result=void  \
14// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
15// RUN: FileCheck %s
16
17#SV = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
18#DV = #sparse_tensor.encoding<{ dimLevelType = [ "dense"      ] }>
19
20#trait_reduction = {
21  indexing_maps = [
22    affine_map<(i) -> (i)>,  // a
23    affine_map<(i) -> ()>    // x (scalar out)
24  ],
25  iterator_types = ["reduction"],
26  doc = "x += OPER_i a(i)"
27}
28
29// An example of vector reductions.
30module {
31
32  func @sum_reduction_i32(%arga: tensor<32xi32, #SV>,
33                          %argx: tensor<i32>) -> tensor<i32> {
34    %0 = linalg.generic #trait_reduction
35      ins(%arga: tensor<32xi32, #SV>)
36      outs(%argx: tensor<i32>) {
37        ^bb(%a: i32, %x: i32):
38          %0 = arith.addi %x, %a : i32
39          linalg.yield %0 : i32
40    } -> tensor<i32>
41    return %0 : tensor<i32>
42  }
43
44  func @sum_reduction_f32(%arga: tensor<32xf32, #SV>,
45                          %argx: tensor<f32>) -> tensor<f32> {
46    %0 = linalg.generic #trait_reduction
47      ins(%arga: tensor<32xf32, #SV>)
48      outs(%argx: tensor<f32>) {
49        ^bb(%a: f32, %x: f32):
50          %0 = arith.addf %x, %a : f32
51          linalg.yield %0 : f32
52    } -> tensor<f32>
53    return %0 : tensor<f32>
54  }
55
56  func @prod_reduction_i32(%arga: tensor<32xi32, #DV>,
57                           %argx: tensor<i32>) -> tensor<i32> {
58    %0 = linalg.generic #trait_reduction
59      ins(%arga: tensor<32xi32, #DV>)
60      outs(%argx: tensor<i32>) {
61        ^bb(%a: i32, %x: i32):
62          %0 = arith.muli %x, %a : i32
63          linalg.yield %0 : i32
64    } -> tensor<i32>
65    return %0 : tensor<i32>
66  }
67
68  func @prod_reduction_f32(%arga: tensor<32xf32, #DV>,
69                           %argx: tensor<f32>) -> tensor<f32> {
70    %0 = linalg.generic #trait_reduction
71      ins(%arga: tensor<32xf32, #DV>)
72      outs(%argx: tensor<f32>) {
73        ^bb(%a: f32, %x: f32):
74          %0 = arith.mulf %x, %a : f32
75          linalg.yield %0 : f32
76    } -> tensor<f32>
77    return %0 : tensor<f32>
78  }
79
80  func @and_reduction_i32(%arga: tensor<32xi32, #DV>,
81                          %argx: tensor<i32>) -> tensor<i32> {
82    %0 = linalg.generic #trait_reduction
83      ins(%arga: tensor<32xi32, #DV>)
84      outs(%argx: tensor<i32>) {
85        ^bb(%a: i32, %x: i32):
86          %0 = arith.andi %x, %a : i32
87          linalg.yield %0 : i32
88    } -> tensor<i32>
89    return %0 : tensor<i32>
90  }
91
92  func @or_reduction_i32(%arga: tensor<32xi32, #SV>,
93                         %argx: tensor<i32>) -> tensor<i32> {
94    %0 = linalg.generic #trait_reduction
95      ins(%arga: tensor<32xi32, #SV>)
96      outs(%argx: tensor<i32>) {
97        ^bb(%a: i32, %x: i32):
98          %0 = arith.ori %x, %a : i32
99          linalg.yield %0 : i32
100    } -> tensor<i32>
101    return %0 : tensor<i32>
102  }
103
104  func @xor_reduction_i32(%arga: tensor<32xi32, #SV>,
105                          %argx: tensor<i32>) -> tensor<i32> {
106    %0 = linalg.generic #trait_reduction
107      ins(%arga: tensor<32xi32, #SV>)
108      outs(%argx: tensor<i32>) {
109        ^bb(%a: i32, %x: i32):
110          %0 = arith.xori %x, %a : i32
111          linalg.yield %0 : i32
112    } -> tensor<i32>
113    return %0 : tensor<i32>
114  }
115
116  func @dump_i32(%arg0 : memref<i32>) {
117    %v = memref.load %arg0[] : memref<i32>
118    vector.print %v : i32
119    return
120  }
121
122  func @dump_f32(%arg0 : memref<f32>) {
123    %v = memref.load %arg0[] : memref<f32>
124    vector.print %v : f32
125    return
126  }
127
128  func @entry() {
129    %ri = arith.constant dense< 7   > : tensor<i32>
130    %rf = arith.constant dense< 2.0 > : tensor<f32>
131
132    %c_0_i32 = arith.constant dense<[
133      0, 2, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 4, 0, 0, 0,
134      0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 9, 0
135    ]> : tensor<32xi32>
136
137    %c_0_f32 = arith.constant dense<[
138      0.0, 1.0, 0.0, 0.0, 4.0, 0.0, 0.0, 0.0,
139      0.0, 0.0, 3.0, 0.0, 0.0, 0.0, 0.0, 0.0,
140      0.0, 0.0, 0.0, 0.0, 2.5, 0.0, 0.0, 0.0,
141      2.0, 0.0, 0.0, 0.0, 0.0, 4.0, 0.0, 9.0
142    ]> : tensor<32xf32>
143
144    %c_1_i32 = arith.constant dense<[
145      1, 1, 7, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
146      1, 1, 1, 1, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1, 7, 3
147    ]> : tensor<32xi32>
148
149    %c_1_f32 = arith.constant dense<[
150      1.0, 1.0, 1.0, 3.5, 1.0, 1.0, 1.0, 1.0,
151      1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0,
152      1.0, 1.0, 1.0, 1.0, 3.0, 1.0, 1.0, 1.0,
153      1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 4.0
154    ]> : tensor<32xf32>
155
156    // Convert constants to annotated tensors.
157    %sparse_input_i32 = sparse_tensor.convert %c_0_i32
158      : tensor<32xi32> to tensor<32xi32, #SV>
159    %sparse_input_f32 = sparse_tensor.convert %c_0_f32
160      : tensor<32xf32> to tensor<32xf32, #SV>
161    %dense_input_i32 = sparse_tensor.convert %c_1_i32
162      : tensor<32xi32> to tensor<32xi32, #DV>
163    %dense_input_f32 = sparse_tensor.convert %c_1_f32
164      : tensor<32xf32> to tensor<32xf32, #DV>
165
166    // Call the kernels.
167    %0 = call @sum_reduction_i32(%sparse_input_i32, %ri)
168       : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
169    %1 = call @sum_reduction_f32(%sparse_input_f32, %rf)
170       : (tensor<32xf32, #SV>, tensor<f32>) -> tensor<f32>
171    %2 = call @prod_reduction_i32(%dense_input_i32, %ri)
172       : (tensor<32xi32, #DV>, tensor<i32>) -> tensor<i32>
173    %3 = call @prod_reduction_f32(%dense_input_f32, %rf)
174       : (tensor<32xf32, #DV>, tensor<f32>) -> tensor<f32>
175    %4 = call @and_reduction_i32(%dense_input_i32, %ri)
176       : (tensor<32xi32, #DV>, tensor<i32>) -> tensor<i32>
177    %5 = call @or_reduction_i32(%sparse_input_i32, %ri)
178       : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
179    %6 = call @xor_reduction_i32(%sparse_input_i32, %ri)
180       : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
181
182    // Verify results.
183    //
184    // CHECK: 26
185    // CHECK: 27.5
186    // CHECK: 3087
187    // CHECK: 168
188    // CHECK: 1
189    // CHECK: 15
190    // CHECK: 10
191    //
192    %m0 = bufferization.to_memref %0 : memref<i32>
193    call @dump_i32(%m0) : (memref<i32>) -> ()
194    %m1 = bufferization.to_memref %1 : memref<f32>
195    call @dump_f32(%m1) : (memref<f32>) -> ()
196    %m2 = bufferization.to_memref %2 : memref<i32>
197    call @dump_i32(%m2) : (memref<i32>) -> ()
198    %m3 = bufferization.to_memref %3 : memref<f32>
199    call @dump_f32(%m3) : (memref<f32>) -> ()
200    %m4 = bufferization.to_memref %4 : memref<i32>
201    call @dump_i32(%m4) : (memref<i32>) -> ()
202    %m5 = bufferization.to_memref %5 : memref<i32>
203    call @dump_i32(%m5) : (memref<i32>) -> ()
204    %m6 = bufferization.to_memref %6 : memref<i32>
205    call @dump_i32(%m6) : (memref<i32>) -> ()
206
207    // Release the resources.
208    sparse_tensor.release %sparse_input_i32 : tensor<32xi32, #SV>
209    sparse_tensor.release %sparse_input_f32 : tensor<32xf32, #SV>
210    sparse_tensor.release %dense_input_i32  : tensor<32xi32, #DV>
211    sparse_tensor.release %dense_input_f32  : tensor<32xf32, #DV>
212    memref.dealloc %m0 : memref<i32>
213    memref.dealloc %m1 : memref<f32>
214    memref.dealloc %m2 : memref<i32>
215    memref.dealloc %m3 : memref<f32>
216    memref.dealloc %m4 : memref<i32>
217    memref.dealloc %m5 : memref<i32>
218    memref.dealloc %m6 : memref<i32>
219
220    return
221  }
222}
223