xref: /llvm-project/mlir/test/Dialect/SparseTensor/vectorize_reduction.mlir (revision ced2fc7819d5ddea616ec330f18e08ff284c1868)
1// RUN: mlir-opt %s --sparse-reinterpret-map -sparsification -cse -sparse-vectorization="vl=8" -cse -split-input-file | \
2// RUN:   FileCheck %s --check-prefix=CHECK-ON
3// RUN: mlir-opt %s --sparse-reinterpret-map -sparsification -cse -split-input-file | \
4// RUN:   FileCheck %s --check-prefix=CHECK-OFF
5
6// -----
7
8// Check that we vectorize reductions with ori.
9
10// CHECK-ON-LABEL:   func.func @sparse_reduction_ori(
11// CHECK-ON-SAME:      %[[VAL_0:.*]]: tensor<i13>,
12// CHECK-ON-SAME:      %[[VAL_1:.*]]: tensor<?xi13, #sparse{{[0-9]*}}>) -> tensor<i13> {
13// CHECK-ON-DAG:       %[[VAL_2:.*]] = arith.constant 8 : index
14// CHECK-ON-DAG:       %[[VAL_3:.*]] = arith.constant dense<0> : vector<8xi13>
15// CHECK-ON-DAG:       %[[VAL_4:.*]] = arith.constant 0 : index
16// CHECK-ON-DAG:       %[[VAL_5:.*]] = arith.constant 1 : index
17// CHECK-ON-DAG:       %[[VAL_6:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xi13, #sparse{{[0-9]*}}> to memref<?xindex>
18// CHECK-ON-DAG:       %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xi13, #sparse{{[0-9]*}}> to memref<?xi13>
19// CHECK-ON-DAG:       %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_0]] : tensor<i13> to memref<i13>
20// CHECK-ON:           %[[VAL_9:.*]] = memref.load %[[VAL_8]][] : memref<i13>
21// CHECK-ON:           %[[VAL_10:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
22// CHECK-ON:           %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
23// CHECK-ON:           %[[VAL_12:.*]] = vector.broadcast %[[VAL_9]] : i13 to vector<8xi13>
24// CHECK-ON:           %[[VAL_13:.*]] = scf.for %[[VAL_14:.*]] = %[[VAL_10]] to %[[VAL_11]] step %[[VAL_2]] iter_args(%[[VAL_15:.*]] = %[[VAL_12]]) -> (vector<8xi13>) {
25// CHECK-ON:             %[[VAL_16:.*]] = affine.min #map(%[[VAL_11]], %[[VAL_14]]){{\[}}%[[VAL_2]]]
26// CHECK-ON:             %[[VAL_17:.*]] = vector.create_mask %[[VAL_16]] : vector<8xi1>
27// CHECK-ON:             %[[VAL_18:.*]] = vector.maskedload %[[VAL_7]]{{\[}}%[[VAL_14]]], %[[VAL_17]], %[[VAL_3]] : memref<?xi13>, vector<8xi1>, vector<8xi13> into vector<8xi13>
28// CHECK-ON:             %[[VAL_19:.*]] = arith.ori %[[VAL_15]], %[[VAL_18]] : vector<8xi13>
29// CHECK-ON:             %[[VAL_20:.*]] = arith.select %[[VAL_17]], %[[VAL_19]], %[[VAL_15]] : vector<8xi1>, vector<8xi13>
30// CHECK-ON:             scf.yield %[[VAL_20]] : vector<8xi13>
31// CHECK-ON:           } {"Emitted from" = "linalg.generic"}
32// CHECK-ON:           %[[VAL_21:.*]] = vector.reduction <or>, %[[VAL_22:.*]] : vector<8xi13> into i13
33// CHECK-ON:           memref.store %[[VAL_21]], %[[VAL_8]][] : memref<i13>
34// CHECK-ON:           %[[VAL_23:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<i13>
35// CHECK-ON:           return %[[VAL_23]] : tensor<i13>
36// CHECK-ON:         }
37//
38// CHECK-OFF-LABEL:   func.func @sparse_reduction_ori(
39// CHECK-OFF-SAME:      %[[VAL_0:.*]]: tensor<i13>,
40// CHECK-OFF-SAME:      %[[VAL_1:.*]]: tensor<?xi13, #sparse{{[0-9]*}}>) -> tensor<i13> {
41// CHECK-OFF-DAG:       %[[VAL_2:.*]] = arith.constant 0 : index
42// CHECK-OFF-DAG:       %[[VAL_3:.*]] = arith.constant 1 : index
43// CHECK-OFF-DAG:       %[[VAL_4:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xi13, #sparse{{[0-9]*}}> to memref<?xindex>
44// CHECK-OFF-DAG:       %[[VAL_5:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xi13, #sparse{{[0-9]*}}> to memref<?xi13>
45// CHECK-OFF-DAG:       %[[VAL_6:.*]] = bufferization.to_memref %[[VAL_0]] : tensor<i13> to memref<i13>
46// CHECK-OFF:           %[[VAL_7:.*]] = memref.load %[[VAL_6]][] : memref<i13>
47// CHECK-OFF:           %[[VAL_8:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_2]]] : memref<?xindex>
48// CHECK-OFF:           %[[VAL_9:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_3]]] : memref<?xindex>
49// CHECK-OFF:           %[[VAL_10:.*]] = scf.for %[[VAL_11:.*]] = %[[VAL_8]] to %[[VAL_9]] step %[[VAL_3]] iter_args(%[[VAL_12:.*]] = %[[VAL_7]]) -> (i13) {
50// CHECK-OFF:             %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_11]]] : memref<?xi13>
51// CHECK-OFF:             %[[VAL_14:.*]] = arith.ori %[[VAL_12]], %[[VAL_13]] : i13
52// CHECK-OFF:             scf.yield %[[VAL_14]] : i13
53// CHECK-OFF:           } {"Emitted from" = "linalg.generic"}
54// CHECK-OFF:           memref.store %[[VAL_15:.*]], %[[VAL_6]][] : memref<i13>
55// CHECK-OFF:           %[[VAL_16:.*]] = bufferization.to_tensor %[[VAL_6]] : memref<i13>
56// CHECK-OFF:           return %[[VAL_16]] : tensor<i13>
57// CHECK-OFF:         }
58#SparseVector = #sparse_tensor.encoding<{map = (d0) -> (d0 : compressed)}>
59
60#trait = {
61  indexing_maps = [
62    affine_map<(i) -> (i)>,  // a (in)
63    affine_map<(i) -> ()>    // x (out)
64  ],
65  iterator_types = ["reduction"]
66}
67
68func.func @sparse_reduction_ori(%argx: tensor<i13>,
69                                %arga: tensor<?xi13, #SparseVector>)
70 -> tensor<i13> {
71  %0 = linalg.generic #trait
72     ins(%arga: tensor<?xi13, #SparseVector>)
73      outs(%argx: tensor<i13>) {
74      ^bb(%a: i13, %x: i13):
75        %t = arith.ori %x, %a: i13
76        linalg.yield %t : i13
77  } -> tensor<i13>
78  return %0 : tensor<i13>
79}
80
81// -----
82
83// Same test as sparse_reduction_ori except that the accumulator is on the
84// rhs of the operation. This checks that we can recognize a reduction
85// irrespective to where the accumulator appears on commutative operations.
86
87// CHECK-ON-LABEL:   func.func @sparse_reduction_ori_accumulator_on_rhs(
88// CHECK-ON-SAME:      %[[VAL_0:.*]]: tensor<i13>,
89// CHECK-ON-SAME:      %[[VAL_1:.*]]: tensor<?xi13, #sparse{{[0-9]*}}>) -> tensor<i13> {
90// CHECK-ON-DAG:       %[[VAL_2:.*]] = arith.constant 8 : index
91// CHECK-ON-DAG:       %[[VAL_3:.*]] = arith.constant dense<0> : vector<8xi13>
92// CHECK-ON-DAG:       %[[VAL_4:.*]] = arith.constant 0 : index
93// CHECK-ON-DAG:       %[[VAL_5:.*]] = arith.constant 1 : index
94// CHECK-ON-DAG:       %[[VAL_6:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xi13, #sparse{{[0-9]*}}> to memref<?xindex>
95// CHECK-ON-DAG:       %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xi13, #sparse{{[0-9]*}}> to memref<?xi13>
96// CHECK-ON-DAG:       %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_0]] : tensor<i13> to memref<i13>
97// CHECK-ON:           %[[VAL_9:.*]] = memref.load %[[VAL_8]][] : memref<i13>
98// CHECK-ON:           %[[VAL_10:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
99// CHECK-ON:           %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
100// CHECK-ON:           %[[VAL_12:.*]] = vector.broadcast %[[VAL_9]] : i13 to vector<8xi13>
101// CHECK-ON:           %[[VAL_13:.*]] = scf.for %[[VAL_14:.*]] = %[[VAL_10]] to %[[VAL_11]] step %[[VAL_2]] iter_args(%[[VAL_15:.*]] = %[[VAL_12]]) -> (vector<8xi13>) {
102// CHECK-ON:             %[[VAL_16:.*]] = affine.min #map(%[[VAL_11]], %[[VAL_14]]){{\[}}%[[VAL_2]]]
103// CHECK-ON:             %[[VAL_17:.*]] = vector.create_mask %[[VAL_16]] : vector<8xi1>
104// CHECK-ON:             %[[VAL_18:.*]] = vector.maskedload %[[VAL_7]]{{\[}}%[[VAL_14]]], %[[VAL_17]], %[[VAL_3]] : memref<?xi13>, vector<8xi1>, vector<8xi13> into vector<8xi13>
105// CHECK-ON:             %[[VAL_19:.*]] = arith.ori %[[VAL_18]], %[[VAL_15]] : vector<8xi13>
106// CHECK-ON:             %[[VAL_20:.*]] = arith.select %[[VAL_17]], %[[VAL_19]], %[[VAL_15]] : vector<8xi1>, vector<8xi13>
107// CHECK-ON:             scf.yield %[[VAL_20]] : vector<8xi13>
108// CHECK-ON:           } {"Emitted from" = "linalg.generic"}
109// CHECK-ON:           %[[VAL_21:.*]] = vector.reduction <or>, %[[VAL_22:.*]] : vector<8xi13> into i13
110// CHECK-ON:           memref.store %[[VAL_21]], %[[VAL_8]][] : memref<i13>
111// CHECK-ON:           %[[VAL_23:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<i13>
112// CHECK-ON:           return %[[VAL_23]] : tensor<i13>
113// CHECK-ON:         }
114//
115// CHECK-OFF-LABEL:   func.func @sparse_reduction_ori_accumulator_on_rhs(
116// CHECK-OFF-SAME:      %[[VAL_0:.*]]: tensor<i13>,
117// CHECK-OFF-SAME:      %[[VAL_1:.*]]: tensor<?xi13, #sparse{{[0-9]*}}>) -> tensor<i13> {
118// CHECK-OFF-DAG:       %[[VAL_2:.*]] = arith.constant 0 : index
119// CHECK-OFF-DAG:       %[[VAL_3:.*]] = arith.constant 1 : index
120// CHECK-OFF-DAG:       %[[VAL_4:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xi13, #sparse{{[0-9]*}}> to memref<?xindex>
121// CHECK-OFF-DAG:       %[[VAL_5:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xi13, #sparse{{[0-9]*}}> to memref<?xi13>
122// CHECK-OFF-DAG:       %[[VAL_6:.*]] = bufferization.to_memref %[[VAL_0]] : tensor<i13> to memref<i13>
123// CHECK-OFF:           %[[VAL_7:.*]] = memref.load %[[VAL_6]][] : memref<i13>
124// CHECK-OFF:           %[[VAL_8:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_2]]] : memref<?xindex>
125// CHECK-OFF:           %[[VAL_9:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_3]]] : memref<?xindex>
126// CHECK-OFF:           %[[VAL_10:.*]] = scf.for %[[VAL_11:.*]] = %[[VAL_8]] to %[[VAL_9]] step %[[VAL_3]] iter_args(%[[VAL_12:.*]] = %[[VAL_7]]) -> (i13) {
127// CHECK-OFF:             %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_11]]] : memref<?xi13>
128// CHECK-OFF:             %[[VAL_14:.*]] = arith.ori %[[VAL_13]], %[[VAL_12]] : i13
129// CHECK-OFF:             scf.yield %[[VAL_14]] : i13
130// CHECK-OFF:           } {"Emitted from" = "linalg.generic"}
131// CHECK-OFF:           memref.store %[[VAL_15:.*]], %[[VAL_6]][] : memref<i13>
132// CHECK-OFF:           %[[VAL_16:.*]] = bufferization.to_tensor %[[VAL_6]] : memref<i13>
133// CHECK-OFF:           return %[[VAL_16]] : tensor<i13>
134// CHECK-OFF:         }
135#SparseVector = #sparse_tensor.encoding<{map = (d0) -> (d0 : compressed)}>
136
137#trait = {
138  indexing_maps = [
139    affine_map<(i) -> (i)>,  // a (in)
140    affine_map<(i) -> ()>    // x (out)
141  ],
142  iterator_types = ["reduction"]
143}
144
145func.func @sparse_reduction_ori_accumulator_on_rhs(%argx: tensor<i13>,
146                                                   %arga: tensor<?xi13, #SparseVector>)
147 -> tensor<i13> {
148  %0 = linalg.generic #trait
149     ins(%arga: tensor<?xi13, #SparseVector>)
150      outs(%argx: tensor<i13>) {
151      ^bb(%a: i13, %x: i13):
152        %t = arith.ori %a, %x: i13
153        linalg.yield %t : i13
154  } -> tensor<i13>
155  return %0 : tensor<i13>
156}
157
158// -----
159
160// Check that we vectorize reductions with subi.
161//
162// CHECK-ON-LABEL:   func.func @sparse_reduction_subi(
163// CHECK-ON-SAME:      %[[VAL_0:.*]]: tensor<i32>,
164// CHECK-ON-SAME:      %[[VAL_1:.*]]: tensor<?xi32, #sparse{{[0-9]*}}>) -> tensor<i32> {
165// CHECK-ON-DAG:       %[[VAL_2:.*]] = arith.constant 8 : index
166// CHECK-ON-DAG:       %[[VAL_3:.*]] = arith.constant 0 : index
167// CHECK-ON-DAG:       %[[VAL_4:.*]] = arith.constant dense<0> : vector<8xi32>
168// CHECK-ON-DAG:       %[[VAL_5:.*]] = arith.constant 1 : index
169// CHECK-ON-DAG:       %[[VAL_6:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xindex>
170// CHECK-ON-DAG:       %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xi32>
171// CHECK-ON-DAG:       %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_0]] : tensor<i32> to memref<i32>
172// CHECK-ON:           %[[VAL_9:.*]] = memref.load %[[VAL_8]][] : memref<i32>
173// CHECK-ON:           %[[VAL_10:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_3]]] : memref<?xindex>
174// CHECK-ON:           %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
175// CHECK-ON:           %[[VAL_12:.*]] = vector.insertelement %[[VAL_9]], %[[VAL_4]]{{\[}}%[[VAL_3]] : index] : vector<8xi32>
176// CHECK-ON:           %[[VAL_13:.*]] = scf.for %[[VAL_14:.*]] = %[[VAL_10]] to %[[VAL_11]] step %[[VAL_2]] iter_args(%[[VAL_15:.*]] = %[[VAL_12]]) -> (vector<8xi32>) {
177// CHECK-ON:             %[[VAL_16:.*]] = affine.min #map(%[[VAL_11]], %[[VAL_14]]){{\[}}%[[VAL_2]]]
178// CHECK-ON:             %[[VAL_17:.*]] = vector.create_mask %[[VAL_16]] : vector<8xi1>
179// CHECK-ON:             %[[VAL_18:.*]] = vector.maskedload %[[VAL_7]]{{\[}}%[[VAL_14]]], %[[VAL_17]], %[[VAL_4]] : memref<?xi32>, vector<8xi1>, vector<8xi32> into vector<8xi32>
180// CHECK-ON:             %[[VAL_19:.*]] = arith.subi %[[VAL_15]], %[[VAL_18]] : vector<8xi32>
181// CHECK-ON:             %[[VAL_20:.*]] = arith.select %[[VAL_17]], %[[VAL_19]], %[[VAL_15]] : vector<8xi1>, vector<8xi32>
182// CHECK-ON:             scf.yield %[[VAL_20]] : vector<8xi32>
183// CHECK-ON:           } {"Emitted from" = "linalg.generic"}
184// CHECK-ON:           %[[VAL_21:.*]] = vector.reduction <add>, %[[VAL_22:.*]] : vector<8xi32> into i32
185// CHECK-ON:           memref.store %[[VAL_21]], %[[VAL_8]][] : memref<i32>
186// CHECK-ON:           %[[VAL_23:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<i32>
187// CHECK-ON:           return %[[VAL_23]] : tensor<i32>
188// CHECK-ON:         }
189//
190// CHECK-OFF-LABEL:   func.func @sparse_reduction_subi(
191// CHECK-OFF-SAME:      %[[VAL_0:.*]]: tensor<i32>,
192// CHECK-OFF-SAME:      %[[VAL_1:.*]]: tensor<?xi32, #sparse{{[0-9]*}}>) -> tensor<i32> {
193// CHECK-OFF-DAG:       %[[VAL_2:.*]] = arith.constant 0 : index
194// CHECK-OFF-DAG:       %[[VAL_3:.*]] = arith.constant 1 : index
195// CHECK-OFF-DAG:       %[[VAL_4:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xindex>
196// CHECK-OFF-DAG:       %[[VAL_5:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xi32>
197// CHECK-OFF-DAG:       %[[VAL_6:.*]] = bufferization.to_memref %[[VAL_0]] : tensor<i32> to memref<i32>
198// CHECK-OFF:           %[[VAL_7:.*]] = memref.load %[[VAL_6]][] : memref<i32>
199// CHECK-OFF:           %[[VAL_8:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_2]]] : memref<?xindex>
200// CHECK-OFF:           %[[VAL_9:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_3]]] : memref<?xindex>
201// CHECK-OFF:           %[[VAL_10:.*]] = scf.for %[[VAL_11:.*]] = %[[VAL_8]] to %[[VAL_9]] step %[[VAL_3]] iter_args(%[[VAL_12:.*]] = %[[VAL_7]]) -> (i32) {
202// CHECK-OFF:             %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_11]]] : memref<?xi32>
203// CHECK-OFF:             %[[VAL_14:.*]] = arith.subi %[[VAL_12]], %[[VAL_13]] : i32
204// CHECK-OFF:             scf.yield %[[VAL_14]] : i32
205// CHECK-OFF:           } {"Emitted from" = "linalg.generic"}
206// CHECK-OFF:           memref.store %[[VAL_15:.*]], %[[VAL_6]][] : memref<i32>
207// CHECK-OFF:           %[[VAL_16:.*]] = bufferization.to_tensor %[[VAL_6]] : memref<i32>
208// CHECK-OFF:           return %[[VAL_16]] : tensor<i32>
209// CHECK-OFF:         }
210#SparseVector = #sparse_tensor.encoding<{map = (d0) -> (d0 : compressed)}>
211
212#trait = {
213  indexing_maps = [
214    affine_map<(i) -> (i)>,  // a (in)
215    affine_map<(i) -> ()>    // x (out)
216  ],
217  iterator_types = ["reduction"]
218}
219
220func.func @sparse_reduction_subi(%argx: tensor<i32>,
221                                 %arga: tensor<?xi32, #SparseVector>)
222 -> tensor<i32> {
223  %0 = linalg.generic #trait
224     ins(%arga: tensor<?xi32, #SparseVector>)
225      outs(%argx: tensor<i32>) {
226      ^bb(%a: i32, %x: i32):
227        %t = arith.subi %x, %a: i32
228        linalg.yield %t : i32
229  } -> tensor<i32>
230  return %0 : tensor<i32>
231}
232
233// -----
234
235// Check that we vectorize reductions with xor.
236
237// CHECK-ON-LABEL: func.func @sparse_reduction_xor(
238// CHECK-ON-SAME: %[[VAL_0:.*]]: tensor<i32>,
239// CHECK-ON-SAME: %[[VAL_1:.*]]: tensor<?xi32, #sparse{{[0-9]*}}>) -> tensor<i32> {
240// CHECK-ON-DAG:  %[[VAL_2:.*]] = arith.constant 8 : index
241// CHECK-ON-DAG:  %[[VAL_3:.*]] = arith.constant dense<0> : vector<8xi32>
242// CHECK-ON-DAG:  %[[VAL_4:.*]] = arith.constant 0 : index
243// CHECK-ON-DAG:  %[[VAL_5:.*]] = arith.constant 1 : index
244// CHECK-ON-DAG:  %[[VAL_6:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xindex>
245// CHECK-ON-DAG:  %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xi32>
246// CHECK-ON-DAG:  %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_0]] : tensor<i32> to memref<i32>
247// CHECK-ON:  %[[VAL_9:.*]] = memref.load %[[VAL_8]][] : memref<i32>
248// CHECK-ON:  %[[VAL_10:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
249// CHECK-ON:  %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
250// CHECK-ON:  %[[VAL_12:.*]] = vector.insertelement %[[VAL_9]], %[[VAL_3]]{{\[}}%[[VAL_4]] : index] : vector<8xi32>
251// CHECK-ON:  %[[VAL_13:.*]] = scf.for %[[VAL_14:.*]] = %[[VAL_10]] to %[[VAL_11]] step %[[VAL_2]] iter_args(%[[VAL_15:.*]] = %[[VAL_12]]) -> (vector<8xi32>) {
252// CHECK-ON:    %[[VAL_16:.*]] = affine.min #map(%[[VAL_11]], %[[VAL_14]]){{\[}}%[[VAL_2]]]
253// CHECK-ON:    %[[VAL_17:.*]] = vector.create_mask %[[VAL_16]] : vector<8xi1>
254// CHECK-ON:    %[[VAL_18:.*]] = vector.maskedload %[[VAL_7]]{{\[}}%[[VAL_14]]], %[[VAL_17]], %[[VAL_3]] : memref<?xi32>, vector<8xi1>, vector<8xi32> into vector<8xi32>
255// CHECK-ON:    %[[VAL_19:.*]] = arith.xori %[[VAL_15]], %[[VAL_18]] : vector<8xi32>
256// CHECK-ON:    %[[VAL_20:.*]] = arith.select %[[VAL_17]], %[[VAL_19]], %[[VAL_15]] : vector<8xi1>, vector<8xi32>
257// CHECK-ON:    scf.yield %[[VAL_20]] : vector<8xi32>
258// CHECK-ON:  } {"Emitted from" = "linalg.generic"}
259// CHECK-ON:  %[[VAL_21:.*]] = vector.reduction <xor>, %[[VAL_22:.*]] : vector<8xi32> into i32
260// CHECK-ON:  memref.store %[[VAL_21]], %[[VAL_8]][] : memref<i32>
261// CHECK-ON:  %[[VAL_23:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<i32>
262// CHECK-ON:  return %[[VAL_23]] : tensor<i32>
263// CHECK-ON: }
264//
265// CHECK-OFF-LABEL: func.func @sparse_reduction_xor(
266// CHECK-OFF-SAME:  %[[VAL_0:.*]]: tensor<i32>,
267// CHECK-OFF-SAME:  %[[VAL_1:.*]]: tensor<?xi32, #sparse{{[0-9]*}}>) -> tensor<i32> {
268// CHECK-OFF-DAG:   %[[VAL_2:.*]] = arith.constant 0 : index
269// CHECK-OFF-DAG:   %[[VAL_3:.*]] = arith.constant 1 : index
270// CHECK-OFF-DAG:   %[[VAL_4:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xindex>
271// CHECK-OFF-DAG:   %[[VAL_5:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xi32>
272// CHECK-OFF-DAG:   %[[VAL_6:.*]] = bufferization.to_memref %[[VAL_0]] : tensor<i32> to memref<i32>
273// CHECK-OFF:   %[[VAL_7:.*]] = memref.load %[[VAL_6]][] : memref<i32>
274// CHECK-OFF:   %[[VAL_8:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_2]]] : memref<?xindex>
275// CHECK-OFF:   %[[VAL_9:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_3]]] : memref<?xindex>
276// CHECK-OFF:   %[[VAL_10:.*]] = scf.for %[[VAL_11:.*]] = %[[VAL_8]] to %[[VAL_9]] step %[[VAL_3]] iter_args(%[[VAL_12:.*]] = %[[VAL_7]]) -> (i32) {
277// CHECK-OFF:     %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_11]]] : memref<?xi32>
278// CHECK-OFF:     %[[VAL_14:.*]] = arith.xori %[[VAL_12]], %[[VAL_13]] : i32
279// CHECK-OFF:     scf.yield %[[VAL_14]] : i32
280// CHECK-OFF:   } {"Emitted from" = "linalg.generic"}
281// CHECK-OFF:   memref.store %[[VAL_15:.*]], %[[VAL_6]][] : memref<i32>
282// CHECK-OFF:   %[[VAL_16:.*]] = bufferization.to_tensor %[[VAL_6]] : memref<i32>
283// CHECK-OFF:   return %[[VAL_16]] : tensor<i32>
284// CHECK-OFF: }
285
286#SparseVector = #sparse_tensor.encoding<{map = (d0) -> (d0 : compressed)}>
287
288#trait = {
289  indexing_maps = [
290    affine_map<(i) -> (i)>,  // a (in)
291    affine_map<(i) -> ()>    // x (out)
292  ],
293  iterator_types = ["reduction"]
294}
295
296func.func @sparse_reduction_xor(%argx: tensor<i32>,
297                             %arga: tensor<?xi32, #SparseVector>)
298 -> tensor<i32> {
299  %0 = linalg.generic #trait
300     ins(%arga: tensor<?xi32, #SparseVector>)
301      outs(%argx: tensor<i32>) {
302      ^bb(%a: i32, %x: i32):
303        %t = arith.xori %x, %a: i32
304        linalg.yield %t : i32
305  } -> tensor<i32>
306  return %0 : tensor<i32>
307}
308
309// -----
310
311// Check that we vectorize reductions with addi.
312
313// CHECK-ON-LABEL: func.func @sparse_reduction_addi(
314// CHECK-ON-SAME:   %[[VAL_0:.*]]: tensor<i32>,
315// CHECK-ON-SAME:   %[[VAL_1:.*]]: tensor<?xi32, #sparse{{[0-9]*}}>) -> tensor<i32> {
316// CHECK-ON-DAG:   %[[VAL_2:.*]] = arith.constant 8 : index
317// CHECK-ON-DAG:   %[[VAL_3:.*]] = arith.constant dense<0> : vector<8xi32>
318// CHECK-ON-DAG:   %[[VAL_4:.*]] = arith.constant 0 : index
319// CHECK-ON-DAG:   %[[VAL_5:.*]] = arith.constant 1 : index
320// CHECK-ON-DAG:   %[[VAL_6:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xindex>
321// CHECK-ON-DAG:   %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xi32>
322// CHECK-ON-DAG:   %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_0]] : tensor<i32> to memref<i32>
323// CHECK-ON:   %[[VAL_9:.*]] = memref.load %[[VAL_8]][] : memref<i32>
324// CHECK-ON:   %[[VAL_10:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
325// CHECK-ON:   %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
326// CHECK-ON:   %[[VAL_12:.*]] = vector.insertelement %[[VAL_9]], %[[VAL_3]]{{\[}}%[[VAL_4]] : index] : vector<8xi32>
327// CHECK-ON:   %[[VAL_13:.*]] = scf.for %[[VAL_14:.*]] = %[[VAL_10]] to %[[VAL_11]] step %[[VAL_2]] iter_args(%[[VAL_15:.*]] = %[[VAL_12]]) -> (vector<8xi32>) {
328// CHECK-ON:     %[[VAL_16:.*]] = affine.min #map(%[[VAL_11]], %[[VAL_14]]){{\[}}%[[VAL_2]]]
329// CHECK-ON:     %[[VAL_17:.*]] = vector.create_mask %[[VAL_16]] : vector<8xi1>
330// CHECK-ON:     %[[VAL_18:.*]] = vector.maskedload %[[VAL_7]]{{\[}}%[[VAL_14]]], %[[VAL_17]], %[[VAL_3]] : memref<?xi32>, vector<8xi1>, vector<8xi32> into vector<8xi32>
331// CHECK-ON:     %[[VAL_19:.*]] = arith.addi %[[VAL_15]], %[[VAL_18]] : vector<8xi32>
332// CHECK-ON:     %[[VAL_20:.*]] = arith.select %[[VAL_17]], %[[VAL_19]], %[[VAL_15]] : vector<8xi1>, vector<8xi32>
333// CHECK-ON:     scf.yield %[[VAL_20]] : vector<8xi32>
334// CHECK-ON:   } {"Emitted from" = "linalg.generic"}
335// CHECK-ON:   %[[VAL_21:.*]] = vector.reduction <add>, %[[VAL_22:.*]] : vector<8xi32> into i32
336// CHECK-ON:   memref.store %[[VAL_21]], %[[VAL_8]][] : memref<i32>
337// CHECK-ON:   %[[VAL_23:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<i32>
338// CHECK-ON:   return %[[VAL_23]] : tensor<i32>
339// CHECK-ON: }
340//
341// CHECK-OFF-LABEL: func.func @sparse_reduction_addi(
342// CHECK-OFF-SAME:   %[[VAL_0:.*]]: tensor<i32>,
343// CHECK-OFF-SAME:   %[[VAL_1:.*]]: tensor<?xi32, #sparse{{[0-9]*}}>) -> tensor<i32> {
344// CHECK-OFF-DAG:   %[[VAL_2:.*]] = arith.constant 0 : index
345// CHECK-OFF-DAG:   %[[VAL_3:.*]] = arith.constant 1 : index
346// CHECK-OFF-DAG:   %[[VAL_4:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xindex>
347// CHECK-OFF-DAG:   %[[VAL_5:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xi32>
348// CHECK-OFF-DAG:   %[[VAL_6:.*]] = bufferization.to_memref %[[VAL_0]] : tensor<i32> to memref<i32>
349// CHECK-OFF:   %[[VAL_7:.*]] = memref.load %[[VAL_6]][] : memref<i32>
350// CHECK-OFF:   %[[VAL_8:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_2]]] : memref<?xindex>
351// CHECK-OFF:   %[[VAL_9:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_3]]] : memref<?xindex>
352// CHECK-OFF:   %[[VAL_10:.*]] = scf.for %[[VAL_11:.*]] = %[[VAL_8]] to %[[VAL_9]] step %[[VAL_3]] iter_args(%[[VAL_12:.*]] = %[[VAL_7]]) -> (i32) {
353// CHECK-OFF:     %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_11]]] : memref<?xi32>
354// CHECK-OFF:     %[[VAL_14:.*]] = arith.addi %[[VAL_12]], %[[VAL_13]] : i32
355// CHECK-OFF:     scf.yield %[[VAL_14]] : i32
356// CHECK-OFF:   } {"Emitted from" = "linalg.generic"}
357// CHECK-OFF:   memref.store %[[VAL_15:.*]], %[[VAL_6]][] : memref<i32>
358// CHECK-OFF:   %[[VAL_16:.*]] = bufferization.to_tensor %[[VAL_6]] : memref<i32>
359// CHECK-OFF:   return %[[VAL_16]] : tensor<i32>
360// CHECK-OFF: }
361
362#SparseVector = #sparse_tensor.encoding<{map = (d0) -> (d0 : compressed)}>
363
364#trait = {
365  indexing_maps = [
366    affine_map<(i) -> (i)>,  // a (in)
367    affine_map<(i) -> ()>    // x (out)
368  ],
369  iterator_types = ["reduction"]
370}
371
372func.func @sparse_reduction_addi(%argx: tensor<i32>,
373                                 %arga: tensor<?xi32, #SparseVector>)
374 -> tensor<i32> {
375  %0 = linalg.generic #trait
376     ins(%arga: tensor<?xi32, #SparseVector>)
377      outs(%argx: tensor<i32>) {
378      ^bb(%a: i32, %x: i32):
379        %t = arith.addi %x, %a: i32
380        linalg.yield %t : i32
381  } -> tensor<i32>
382  return %0 : tensor<i32>
383}
384
385// -----
386
387// Check that we vectorize reductions with subf.
388
389// CHECK-ON-LABEL: func.func @sparse_reduction_subf(
390// CHECK-ON-SAME:   %[[VAL_0:.*]]: tensor<f32>,
391// CHECK-ON-SAME:   %[[VAL_1:.*]]: tensor<?xf32, #sparse{{[0-9]*}}>) -> tensor<f32> {
392// CHECK-ON-DAG:   %[[VAL_2:.*]] = arith.constant 8 : index
393// CHECK-ON-DAG:   %[[VAL_3:.*]] = arith.constant dense<0.000000e+00> : vector<8xf32>
394// CHECK-ON-DAG:   %[[VAL_4:.*]] = arith.constant 0 : index
395// CHECK-ON-DAG:   %[[VAL_5:.*]] = arith.constant 1 : index
396// CHECK-ON-DAG:   %[[VAL_6:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xf32, #sparse{{[0-9]*}}> to memref<?xindex>
397// CHECK-ON-DAG:   %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xf32, #sparse{{[0-9]*}}> to memref<?xf32>
398// CHECK-ON-DAG:   %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_0]] : tensor<f32> to memref<f32>
399// CHECK-ON:   %[[VAL_9:.*]] = memref.load %[[VAL_8]][] : memref<f32>
400// CHECK-ON:   %[[VAL_10:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
401// CHECK-ON:   %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
402// CHECK-ON:   %[[VAL_12:.*]] = vector.insertelement %[[VAL_9]], %[[VAL_3]]{{\[}}%[[VAL_4]] : index] : vector<8xf32>
403// CHECK-ON:   %[[VAL_13:.*]] = scf.for %[[VAL_14:.*]] = %[[VAL_10]] to %[[VAL_11]] step %[[VAL_2]] iter_args(%[[VAL_15:.*]] = %[[VAL_12]]) -> (vector<8xf32>) {
404// CHECK-ON:     %[[VAL_16:.*]] = affine.min #map(%[[VAL_11]], %[[VAL_14]]){{\[}}%[[VAL_2]]]
405// CHECK-ON:     %[[VAL_17:.*]] = vector.create_mask %[[VAL_16]] : vector<8xi1>
406// CHECK-ON:     %[[VAL_18:.*]] = vector.maskedload %[[VAL_7]]{{\[}}%[[VAL_14]]], %[[VAL_17]], %[[VAL_3]] : memref<?xf32>, vector<8xi1>, vector<8xf32> into vector<8xf32>
407// CHECK-ON:     %[[VAL_19:.*]] = arith.subf %[[VAL_15]], %[[VAL_18]] : vector<8xf32>
408// CHECK-ON:     %[[VAL_20:.*]] = arith.select %[[VAL_17]], %[[VAL_19]], %[[VAL_15]] : vector<8xi1>, vector<8xf32>
409// CHECK-ON:     scf.yield %[[VAL_20]] : vector<8xf32>
410// CHECK-ON:   } {"Emitted from" = "linalg.generic"}
411// CHECK-ON:   %[[VAL_21:.*]] = vector.reduction <add>, %[[VAL_22:.*]] : vector<8xf32> into f32
412// CHECK-ON:   memref.store %[[VAL_21]], %[[VAL_8]][] : memref<f32>
413// CHECK-ON:   %[[VAL_23:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<f32>
414// CHECK-ON:   return %[[VAL_23]] : tensor<f32>
415// CHECK-ON: }
416//
417// CHECK-OFF-LABEL: func.func @sparse_reduction_subf(
418// CHECK-OFF-SAME:   %[[VAL_0:.*]]: tensor<f32>,
419// CHECK-OFF-SAME:   %[[VAL_1:.*]]: tensor<?xf32, #sparse{{[0-9]*}}>) -> tensor<f32> {
420// CHECK-OFF-DAG:   %[[VAL_2:.*]] = arith.constant 0 : index
421// CHECK-OFF-DAG:   %[[VAL_3:.*]] = arith.constant 1 : index
422// CHECK-OFF-DAG:   %[[VAL_4:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xf32, #sparse{{[0-9]*}}> to memref<?xindex>
423// CHECK-OFF-DAG:   %[[VAL_5:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xf32, #sparse{{[0-9]*}}> to memref<?xf32>
424// CHECK-OFF-DAG:   %[[VAL_6:.*]] = bufferization.to_memref %[[VAL_0]] : tensor<f32> to memref<f32>
425// CHECK-OFF:   %[[VAL_7:.*]] = memref.load %[[VAL_6]][] : memref<f32>
426// CHECK-OFF:   %[[VAL_8:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_2]]] : memref<?xindex>
427// CHECK-OFF:   %[[VAL_9:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_3]]] : memref<?xindex>
428// CHECK-OFF:   %[[VAL_10:.*]] = scf.for %[[VAL_11:.*]] = %[[VAL_8]] to %[[VAL_9]] step %[[VAL_3]] iter_args(%[[VAL_12:.*]] = %[[VAL_7]]) -> (f32) {
429// CHECK-OFF:     %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_11]]] : memref<?xf32>
430// CHECK-OFF:     %[[VAL_14:.*]] = arith.subf %[[VAL_12]], %[[VAL_13]] : f32
431// CHECK-OFF:     scf.yield %[[VAL_14]] : f32
432// CHECK-OFF:   } {"Emitted from" = "linalg.generic"}
433// CHECK-OFF:   memref.store %[[VAL_15:.*]], %[[VAL_6]][] : memref<f32>
434// CHECK-OFF:   %[[VAL_16:.*]] = bufferization.to_tensor %[[VAL_6]] : memref<f32>
435// CHECK-OFF:   return %[[VAL_16]] : tensor<f32>
436// CHECK-OFF: }
437
438#SparseVector = #sparse_tensor.encoding<{map = (d0) -> (d0 : compressed)}>
439
440#trait = {
441  indexing_maps = [
442    affine_map<(i) -> (i)>,  // a (in)
443    affine_map<(i) -> ()>    // x (out)
444  ],
445  iterator_types = ["reduction"]
446}
447
448func.func @sparse_reduction_subf(%argx: tensor<f32>,
449                             %arga: tensor<?xf32, #SparseVector>)
450 -> tensor<f32> {
451  %0 = linalg.generic #trait
452     ins(%arga: tensor<?xf32, #SparseVector>)
453      outs(%argx: tensor<f32>) {
454      ^bb(%a: f32, %x: f32):
455        %t = arith.subf %x, %a: f32
456        linalg.yield %t : f32
457  } -> tensor<f32>
458  return %0 : tensor<f32>
459}
460
461// -----
462
463// Check that we vectorize reductions with addf.
464
465// CHECK-ON-LABEL: func.func @sparse_reduction_addf(
466// CHECK-ON-SAME:  %[[VAL_0:.*]]: tensor<f32>,
467// CHECK-ON-SAME:  %[[VAL_1:.*]]: tensor<?xf32, #sparse{{[0-9]*}}>) -> tensor<f32> {
468// CHECK-ON-DAG:   %[[VAL_2:.*]] = arith.constant 8 : index
469// CHECK-ON-DAG:   %[[VAL_3:.*]] = arith.constant dense<0.000000e+00> : vector<8xf32>
470// CHECK-ON-DAG:   %[[VAL_4:.*]] = arith.constant 0 : index
471// CHECK-ON-DAG:   %[[VAL_5:.*]] = arith.constant 1 : index
472// CHECK-ON-DAG:   %[[VAL_6:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xf32, #sparse{{[0-9]*}}> to memref<?xindex>
473// CHECK-ON-DAG:   %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xf32, #sparse{{[0-9]*}}> to memref<?xf32>
474// CHECK-ON-DAG:   %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_0]] : tensor<f32> to memref<f32>
475// CHECK-ON:   %[[VAL_9:.*]] = memref.load %[[VAL_8]][] : memref<f32>
476// CHECK-ON:   %[[VAL_10:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
477// CHECK-ON:   %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
478// CHECK-ON:   %[[VAL_12:.*]] = vector.insertelement %[[VAL_9]], %[[VAL_3]]{{\[}}%[[VAL_4]] : index] : vector<8xf32>
479// CHECK-ON:   %[[VAL_13:.*]] = scf.for %[[VAL_14:.*]] = %[[VAL_10]] to %[[VAL_11]] step %[[VAL_2]] iter_args(%[[VAL_15:.*]] = %[[VAL_12]]) -> (vector<8xf32>) {
480// CHECK-ON:     %[[VAL_16:.*]] = affine.min #map(%[[VAL_11]], %[[VAL_14]]){{\[}}%[[VAL_2]]]
481// CHECK-ON:     %[[VAL_17:.*]] = vector.create_mask %[[VAL_16]] : vector<8xi1>
482// CHECK-ON:     %[[VAL_18:.*]] = vector.maskedload %[[VAL_7]]{{\[}}%[[VAL_14]]], %[[VAL_17]], %[[VAL_3]] : memref<?xf32>, vector<8xi1>, vector<8xf32> into vector<8xf32>
483// CHECK-ON:     %[[VAL_19:.*]] = arith.addf %[[VAL_15]], %[[VAL_18]] : vector<8xf32>
484// CHECK-ON:     %[[VAL_20:.*]] = arith.select %[[VAL_17]], %[[VAL_19]], %[[VAL_15]] : vector<8xi1>, vector<8xf32>
485// CHECK-ON:     scf.yield %[[VAL_20]] : vector<8xf32>
486// CHECK-ON:   } {"Emitted from" = "linalg.generic"}
487// CHECK-ON:   %[[VAL_21:.*]] = vector.reduction <add>, %[[VAL_22:.*]] : vector<8xf32> into f32
488// CHECK-ON:   memref.store %[[VAL_21]], %[[VAL_8]][] : memref<f32>
489// CHECK-ON:   %[[VAL_23:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<f32>
490// CHECK-ON:   return %[[VAL_23]] : tensor<f32>
491// CHECK-ON: }
492//
493// CHECK-OFF-LABEL: func.func @sparse_reduction_addf(
494// CHECK-OFF-SAME:  %[[VAL_0:.*]]: tensor<f32>,
495// CHECK-OFF-SAME:  %[[VAL_1:.*]]: tensor<?xf32, #sparse{{[0-9]*}}>) -> tensor<f32> {
496// CHECK-OFF-DAG:   %[[VAL_2:.*]] = arith.constant 0 : index
497// CHECK-OFF-DAG:   %[[VAL_3:.*]] = arith.constant 1 : index
498// CHECK-OFF-DAG:   %[[VAL_4:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xf32, #sparse{{[0-9]*}}> to memref<?xindex>
499// CHECK-OFF-DAG:   %[[VAL_5:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xf32, #sparse{{[0-9]*}}> to memref<?xf32>
500// CHECK-OFF-DAG:   %[[VAL_6:.*]] = bufferization.to_memref %[[VAL_0]] : tensor<f32> to memref<f32>
501// CHECK-OFF:   %[[VAL_7:.*]] = memref.load %[[VAL_6]][] : memref<f32>
502// CHECK-OFF:   %[[VAL_8:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_2]]] : memref<?xindex>
503// CHECK-OFF:   %[[VAL_9:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_3]]] : memref<?xindex>
504// CHECK-OFF:   %[[VAL_10:.*]] = scf.for %[[VAL_11:.*]] = %[[VAL_8]] to %[[VAL_9]] step %[[VAL_3]] iter_args(%[[VAL_12:.*]] = %[[VAL_7]]) -> (f32) {
505// CHECK-OFF:     %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_11]]] : memref<?xf32>
506// CHECK-OFF:     %[[VAL_14:.*]] = arith.addf %[[VAL_12]], %[[VAL_13]] : f32
507// CHECK-OFF:     scf.yield %[[VAL_14]] : f32
508// CHECK-OFF:   } {"Emitted from" = "linalg.generic"}
509// CHECK-OFF:   memref.store %[[VAL_15:.*]], %[[VAL_6]][] : memref<f32>
510// CHECK-OFF:   %[[VAL_16:.*]] = bufferization.to_tensor %[[VAL_6]] : memref<f32>
511// CHECK-OFF:   return %[[VAL_16]] : tensor<f32>
512// CHECK-OFF: }
513
514#SparseVector = #sparse_tensor.encoding<{map = (d0) -> (d0 : compressed)}>
515
516#trait = {
517  indexing_maps = [
518    affine_map<(i) -> (i)>,  // a (in)
519    affine_map<(i) -> ()>    // x (out)
520  ],
521  iterator_types = ["reduction"]
522}
523
524func.func @sparse_reduction_addf(%argx: tensor<f32>,
525                             %arga: tensor<?xf32, #SparseVector>)
526 -> tensor<f32> {
527  %0 = linalg.generic #trait
528     ins(%arga: tensor<?xf32, #SparseVector>)
529      outs(%argx: tensor<f32>) {
530      ^bb(%a: f32, %x: f32):
531        %t = arith.addf %x, %a: f32
532        linalg.yield %t : f32
533  } -> tensor<f32>
534  return %0 : tensor<f32>
535}
536