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