xref: /llvm-project/mlir/test/Dialect/SparseTensor/sparse_int_ops.mlir (revision ced2fc7819d5ddea616ec330f18e08ff284c1868)
1// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
2// RUN: mlir-opt %s --sparse-reinterpret-map -sparsification | FileCheck %s
3
4#SV = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
5
6#trait2 = {
7  indexing_maps = [
8    affine_map<(i) -> (i)>,  // a
9    affine_map<(i) -> (i)>,  // b
10    affine_map<(i) -> (i)>   // x (out)
11  ],
12  iterator_types = ["parallel"],
13  doc = "x(i) = a(i) OP b(i)"
14}
15
16#traitc = {
17  indexing_maps = [
18    affine_map<(i) -> (i)>,  // a
19    affine_map<(i) -> (i)>   // x (out)
20  ],
21  iterator_types = ["parallel"],
22  doc = "x(i) = a(i) OP c"
23}
24
25// CHECK-LABEL:   func @add(
26// CHECK-SAME:              %[[VAL_0:.*]]: tensor<32xi64, #sparse{{[0-9]*}}>,
27// CHECK-SAME:              %[[VAL_1:.*]]: tensor<32xi64>,
28// CHECK-SAME:              %[[VAL_2:.*]]: tensor<32xi64>) -> tensor<32xi64> {
29// CHECK-DAG:           %[[VAL_3:.*]] = arith.constant 32 : index
30// CHECK-DAG:           %[[VAL_4:.*]] = arith.constant 0 : index
31// CHECK-DAG:           %[[VAL_5:.*]] = arith.constant true
32// CHECK-DAG:           %[[VAL_6:.*]] = arith.constant 1 : index
33// CHECK-DAG:           %[[VAL_7:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}>
34// CHECK-DAG:           %[[VAL_8:.*]] = sparse_tensor.coordinates %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}>
35// CHECK-DAG:           %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xi64, #sparse{{[0-9]*}}>
36// CHECK-DAG:           %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_1]] : tensor<32xi64> to memref<32xi64>
37// CHECK-DAG:           %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : tensor<32xi64> to memref<32xi64>
38// CHECK:           %[[VAL_12:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_4]]] : memref<?xindex>
39// CHECK:           %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex>
40// CHECK:           %[[VAL_14:.*]]:2 = scf.while (%[[VAL_15:.*]] = %[[VAL_12]], %[[VAL_16:.*]] = %[[VAL_4]]) : (index, index) -> (index, index) {
41// CHECK:             %[[VAL_17:.*]] = arith.cmpi ult, %[[VAL_15]], %[[VAL_13]] : index
42// CHECK:             scf.condition(%[[VAL_17]]) %[[VAL_15]], %[[VAL_16]] : index, index
43// CHECK:           } do {
44// CHECK:           ^bb0(%[[VAL_18:.*]]: index, %[[VAL_19:.*]]: index):
45// CHECK:             %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_18]]] : memref<?xindex>
46// CHECK:             %[[VAL_21:.*]] = arith.cmpi eq, %[[VAL_20]], %[[VAL_19]] : index
47// CHECK:             scf.if %[[VAL_21]] {
48// CHECK:               %[[VAL_22:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_18]]] : memref<?xi64>
49// CHECK:               %[[VAL_23:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_19]]] : memref<32xi64>
50// CHECK:               %[[VAL_24:.*]] = arith.addi %[[VAL_22]], %[[VAL_23]] : i64
51// CHECK:               memref.store %[[VAL_24]], %[[VAL_11]]{{\[}}%[[VAL_19]]] : memref<32xi64>
52// CHECK:             } else {
53// CHECK:               scf.if %[[VAL_5]] {
54// CHECK:                 %[[VAL_25:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_19]]] : memref<32xi64>
55// CHECK:                 memref.store %[[VAL_25]], %[[VAL_11]]{{\[}}%[[VAL_19]]] : memref<32xi64>
56// CHECK:               } else {
57// CHECK:               }
58// CHECK:             }
59// CHECK:             %[[VAL_26:.*]] = arith.cmpi eq, %[[VAL_20]], %[[VAL_19]] : index
60// CHECK:             %[[VAL_27:.*]] = arith.addi %[[VAL_18]], %[[VAL_6]] : index
61// CHECK:             %[[VAL_28:.*]] = arith.select %[[VAL_26]], %[[VAL_27]], %[[VAL_18]] : index
62// CHECK:             %[[VAL_29:.*]] = arith.addi %[[VAL_19]], %[[VAL_6]] : index
63// CHECK:             scf.yield %[[VAL_28]], %[[VAL_29]] : index, index
64// CHECK:           }
65// CHECK:           scf.for %[[VAL_30:.*]] = %[[VAL_31:.*]]#1 to %[[VAL_3]] step %[[VAL_6]] {
66// CHECK:             %[[VAL_32:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_30]]] : memref<32xi64>
67// CHECK:             memref.store %[[VAL_32]], %[[VAL_11]]{{\[}}%[[VAL_30]]] : memref<32xi64>
68// CHECK:           }
69// CHECK:           %[[VAL_33:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xi64>
70// CHECK:           return %[[VAL_33]] : tensor<32xi64>
71// CHECK:         }
72func.func @add(%arga: tensor<32xi64, #SV>,
73               %argb: tensor<32xi64>,
74               %argx: tensor<32xi64>) -> tensor<32xi64> {
75  %0 = linalg.generic #trait2
76     ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>)
77    outs(%argx: tensor<32xi64>) {
78      ^bb(%a: i64, %b: i64, %x: i64):
79        %0 = arith.addi %a, %b : i64
80        linalg.yield %0 : i64
81  } -> tensor<32xi64>
82  return %0 : tensor<32xi64>
83}
84
85// CHECK-LABEL:   func @sub(
86// CHECK-SAME:              %[[VAL_0:.*]]: tensor<32xi64, #sparse{{[0-9]*}}>,
87// CHECK-SAME:              %[[VAL_1:.*]]: tensor<32xi64>,
88// CHECK-SAME:              %[[VAL_2:.*]]: tensor<32xi64>) -> tensor<32xi64> {
89// CHECK-DAG:           %[[VAL_3:.*]] = arith.constant 32 : index
90// CHECK-DAG:           %[[VAL_4:.*]] = arith.constant 0 : index
91// CHECK-DAG:           %[[VAL_5:.*]] = arith.constant true
92// CHECK-DAG:           %[[VAL_6:.*]] = arith.constant 1 : index
93// CHECK-DAG:           %[[VAL_7:.*]] = arith.constant 0 : i64
94// CHECK-DAG:           %[[VAL_8:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}>
95// CHECK-DAG:           %[[VAL_9:.*]] = sparse_tensor.coordinates %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}>
96// CHECK-DAG:           %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xi64, #sparse{{[0-9]*}}>
97// CHECK-DAG:           %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_1]] : tensor<32xi64> to memref<32xi64>
98// CHECK-DAG:           %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]] : tensor<32xi64> to memref<32xi64>
99// CHECK:           %[[VAL_13:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_4]]] : memref<?xindex>
100// CHECK:           %[[VAL_14:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_6]]] : memref<?xindex>
101// CHECK:           %[[VAL_15:.*]]:2 = scf.while (%[[VAL_16:.*]] = %[[VAL_13]], %[[VAL_17:.*]] = %[[VAL_4]]) : (index, index) -> (index, index) {
102// CHECK:             %[[VAL_18:.*]] = arith.cmpi ult, %[[VAL_16]], %[[VAL_14]] : index
103// CHECK:             scf.condition(%[[VAL_18]]) %[[VAL_16]], %[[VAL_17]] : index, index
104// CHECK:           } do {
105// CHECK:           ^bb0(%[[VAL_19:.*]]: index, %[[VAL_20:.*]]: index):
106// CHECK:             %[[VAL_21:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_19]]] : memref<?xindex>
107// CHECK:             %[[VAL_22:.*]] = arith.cmpi eq, %[[VAL_21]], %[[VAL_20]] : index
108// CHECK:             scf.if %[[VAL_22]] {
109// CHECK:               %[[VAL_23:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_19]]] : memref<?xi64>
110// CHECK:               %[[VAL_24:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_20]]] : memref<32xi64>
111// CHECK:               %[[VAL_25:.*]] = arith.subi %[[VAL_23]], %[[VAL_24]] : i64
112// CHECK:               memref.store %[[VAL_25]], %[[VAL_12]]{{\[}}%[[VAL_20]]] : memref<32xi64>
113// CHECK:             } else {
114// CHECK:               scf.if %[[VAL_5]] {
115// CHECK:                 %[[VAL_26:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_20]]] : memref<32xi64>
116// CHECK:                 %[[VAL_27:.*]] = arith.subi %[[VAL_7]], %[[VAL_26]] : i64
117// CHECK:                 memref.store %[[VAL_27]], %[[VAL_12]]{{\[}}%[[VAL_20]]] : memref<32xi64>
118// CHECK:               } else {
119// CHECK:               }
120// CHECK:             }
121// CHECK:             %[[VAL_28:.*]] = arith.cmpi eq, %[[VAL_21]], %[[VAL_20]] : index
122// CHECK:             %[[VAL_29:.*]] = arith.addi %[[VAL_19]], %[[VAL_6]] : index
123// CHECK:             %[[VAL_30:.*]] = arith.select %[[VAL_28]], %[[VAL_29]], %[[VAL_19]] : index
124// CHECK:             %[[VAL_31:.*]] = arith.addi %[[VAL_20]], %[[VAL_6]] : index
125// CHECK:             scf.yield %[[VAL_30]], %[[VAL_31]] : index, index
126// CHECK:           }
127// CHECK:           scf.for %[[VAL_32:.*]] = %[[VAL_33:.*]]#1 to %[[VAL_3]] step %[[VAL_6]] {
128// CHECK:             %[[VAL_34:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_32]]] : memref<32xi64>
129// CHECK:             %[[VAL_35:.*]] = arith.subi %[[VAL_7]], %[[VAL_34]] : i64
130// CHECK:             memref.store %[[VAL_35]], %[[VAL_12]]{{\[}}%[[VAL_32]]] : memref<32xi64>
131// CHECK:           }
132// CHECK:           %[[VAL_36:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<32xi64>
133// CHECK:           return %[[VAL_36]] : tensor<32xi64>
134// CHECK:         }
135func.func @sub(%arga: tensor<32xi64, #SV>,
136               %argb: tensor<32xi64>,
137               %argx: tensor<32xi64>) -> tensor<32xi64> {
138  %0 = linalg.generic #trait2
139     ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>)
140    outs(%argx: tensor<32xi64>) {
141      ^bb(%a: i64, %b: i64, %x: i64):
142        %0 = arith.subi %a, %b : i64
143        linalg.yield %0 : i64
144  } -> tensor<32xi64>
145  return %0 : tensor<32xi64>
146}
147
148// CHECK-LABEL:   func @mul(
149// CHECK-SAME:              %[[VAL_0:.*]]: tensor<32xi64, #sparse{{[0-9]*}}>,
150// CHECK-SAME:              %[[VAL_1:.*]]: tensor<32xi64>,
151// CHECK-SAME:              %[[VAL_2:.*]]: tensor<32xi64>) -> tensor<32xi64> {
152// CHECK-DAG:           %[[VAL_3:.*]] = arith.constant 0 : index
153// CHECK-DAG:           %[[VAL_4:.*]] = arith.constant 1 : index
154// CHECK-DAG:           %[[VAL_5:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}>
155// CHECK-DAG:           %[[VAL_6:.*]] = sparse_tensor.coordinates %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}>
156// CHECK-DAG:           %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xi64, #sparse{{[0-9]*}}>
157// CHECK-DAG:           %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]] : tensor<32xi64> to memref<32xi64>
158// CHECK-DAG:           %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_2]] : tensor<32xi64> to memref<32xi64>
159// CHECK:           %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
160// CHECK:           %[[VAL_11:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
161// CHECK:           scf.for %[[VAL_12:.*]] = %[[VAL_10]] to %[[VAL_11]] step %[[VAL_4]] {
162// CHECK:             %[[VAL_13:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_12]]] : memref<?xindex>
163// CHECK:             %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_12]]] : memref<?xi64>
164// CHECK:             %[[VAL_15:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_13]]] : memref<32xi64>
165// CHECK:             %[[VAL_16:.*]] = arith.muli %[[VAL_14]], %[[VAL_15]] : i64
166// CHECK:             memref.store %[[VAL_16]], %[[VAL_9]]{{\[}}%[[VAL_13]]] : memref<32xi64>
167// CHECK:           }
168// CHECK:           %[[VAL_17:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<32xi64>
169// CHECK:           return %[[VAL_17]] : tensor<32xi64>
170// CHECK:         }
171func.func @mul(%arga: tensor<32xi64, #SV>,
172               %argb: tensor<32xi64>,
173               %argx: tensor<32xi64>) -> tensor<32xi64> {
174  %0 = linalg.generic #trait2
175     ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>)
176    outs(%argx: tensor<32xi64>) {
177      ^bb(%a: i64, %b: i64, %x: i64):
178        %0 = arith.muli %a, %b : i64
179        linalg.yield %0 : i64
180  } -> tensor<32xi64>
181  return %0 : tensor<32xi64>
182}
183
184// CHECK-LABEL:   func @divsbyc(
185// CHECK-SAME:                  %[[VAL_0:.*]]: tensor<32xi64, #sparse{{[0-9]*}}>,
186// CHECK-SAME:                  %[[VAL_1:.*]]: tensor<32xi64>) -> tensor<32xi64> {
187// CHECK-DAG:           %[[VAL_2:.*]] = arith.constant 2 : i64
188// CHECK-DAG:           %[[VAL_3:.*]] = arith.constant 0 : index
189// CHECK-DAG:           %[[VAL_4:.*]] = arith.constant 1 : index
190// CHECK-DAG:           %[[VAL_5:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}>
191// CHECK-DAG:           %[[VAL_6:.*]] = sparse_tensor.coordinates %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}>
192// CHECK-DAG:           %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xi64, #sparse{{[0-9]*}}>
193// CHECK-DAG:           %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]] : tensor<32xi64> to memref<32xi64>
194// CHECK:           %[[VAL_9:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
195// CHECK:           %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
196// CHECK:           scf.for %[[VAL_11:.*]] = %[[VAL_9]] to %[[VAL_10]] step %[[VAL_4]] {
197// CHECK:             %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xindex>
198// CHECK:             %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_11]]] : memref<?xi64>
199// CHECK:             %[[VAL_14:.*]] = arith.divsi %[[VAL_13]], %[[VAL_2]] : i64
200// CHECK:             memref.store %[[VAL_14]], %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<32xi64>
201// CHECK:           }
202// CHECK:           %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xi64>
203// CHECK:           return %[[VAL_15]] : tensor<32xi64>
204// CHECK:         }
205func.func @divsbyc(%arga: tensor<32xi64, #SV>,
206                   %argx: tensor<32xi64>) -> tensor<32xi64> {
207  %c = arith.constant 2 : i64
208  %0 = linalg.generic #traitc
209     ins(%arga: tensor<32xi64, #SV>)
210    outs(%argx: tensor<32xi64>) {
211      ^bb(%a: i64, %x: i64):
212        %0 = arith.divsi %a, %c : i64
213        linalg.yield %0 : i64
214  } -> tensor<32xi64>
215  return %0 : tensor<32xi64>
216}
217
218// CHECK-LABEL:   func @divubyc(
219// CHECK-SAME:                  %[[VAL_0:.*]]: tensor<32xi64, #sparse{{[0-9]*}}>,
220// CHECK-SAME:                  %[[VAL_1:.*]]: tensor<32xi64>) -> tensor<32xi64> {
221// CHECK-DAG:           %[[VAL_2:.*]] = arith.constant 2 : i64
222// CHECK-DAG:           %[[VAL_3:.*]] = arith.constant 0 : index
223// CHECK-DAG:           %[[VAL_4:.*]] = arith.constant 1 : index
224// CHECK-DAG:           %[[VAL_5:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}>
225// CHECK-DAG:           %[[VAL_6:.*]] = sparse_tensor.coordinates %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}>
226// CHECK-DAG:           %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xi64, #sparse{{[0-9]*}}>
227// CHECK-DAG:           %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]] : tensor<32xi64> to memref<32xi64>
228// CHECK:           %[[VAL_9:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
229// CHECK:           %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
230// CHECK:           scf.for %[[VAL_11:.*]] = %[[VAL_9]] to %[[VAL_10]] step %[[VAL_4]] {
231// CHECK:             %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xindex>
232// CHECK:             %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_11]]] : memref<?xi64>
233// CHECK:             %[[VAL_14:.*]] = arith.divui %[[VAL_13]], %[[VAL_2]] : i64
234// CHECK:             memref.store %[[VAL_14]], %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<32xi64>
235// CHECK:           }
236// CHECK:           %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xi64>
237// CHECK:           return %[[VAL_15]] : tensor<32xi64>
238// CHECK:         }
239func.func @divubyc(%arga: tensor<32xi64, #SV>,
240                   %argx: tensor<32xi64>) -> tensor<32xi64> {
241  %c = arith.constant 2 : i64
242  %0 = linalg.generic #traitc
243     ins(%arga: tensor<32xi64, #SV>)
244    outs(%argx: tensor<32xi64>) {
245      ^bb(%a: i64, %x: i64):
246        %0 = arith.divui %a, %c : i64
247        linalg.yield %0 : i64
248  } -> tensor<32xi64>
249  return %0 : tensor<32xi64>
250}
251
252// CHECK-LABEL:   func @and(
253// CHECK-SAME:              %[[VAL_0:.*]]: tensor<32xi64, #sparse{{[0-9]*}}>,
254// CHECK-SAME:              %[[VAL_1:.*]]: tensor<32xi64>,
255// CHECK-SAME:              %[[VAL_2:.*]]: tensor<32xi64>) -> tensor<32xi64> {
256// CHECK-DAG:           %[[VAL_3:.*]] = arith.constant 0 : index
257// CHECK-DAG:           %[[VAL_4:.*]] = arith.constant 1 : index
258// CHECK-DAG:           %[[VAL_5:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}> to memref<?xindex>
259// CHECK-DAG:           %[[VAL_6:.*]] = sparse_tensor.coordinates %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}> to memref<?xindex>
260// CHECK-DAG:           %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xi64, #sparse{{[0-9]*}}> to memref<?xi64>
261// CHECK-DAG:           %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]] : tensor<32xi64> to memref<32xi64>
262// CHECK-DAG:           %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_2]] : tensor<32xi64> to memref<32xi64>
263// CHECK:           %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
264// CHECK:           %[[VAL_11:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
265// CHECK:           scf.for %[[VAL_12:.*]] = %[[VAL_10]] to %[[VAL_11]] step %[[VAL_4]] {
266// CHECK:             %[[VAL_13:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_12]]] : memref<?xindex>
267// CHECK:             %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_12]]] : memref<?xi64>
268// CHECK:             %[[VAL_15:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_13]]] : memref<32xi64>
269// CHECK:             %[[VAL_16:.*]] = arith.andi %[[VAL_14]], %[[VAL_15]] : i64
270// CHECK:             memref.store %[[VAL_16]], %[[VAL_9]]{{\[}}%[[VAL_13]]] : memref<32xi64>
271// CHECK:           }
272// CHECK:           %[[VAL_17:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<32xi64>
273// CHECK:           return %[[VAL_17]] : tensor<32xi64>
274// CHECK:         }
275func.func @and(%arga: tensor<32xi64, #SV>,
276               %argb: tensor<32xi64>,
277               %argx: tensor<32xi64>) -> tensor<32xi64> {
278  %0 = linalg.generic #trait2
279     ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>)
280    outs(%argx: tensor<32xi64>) {
281      ^bb(%a: i64, %b: i64, %x: i64):
282        %0 = arith.andi %a, %b : i64
283        linalg.yield %0 : i64
284  } -> tensor<32xi64>
285  return %0 : tensor<32xi64>
286}
287
288// CHECK-LABEL:   func @or(
289// CHECK-SAME:             %[[VAL_0:.*]]: tensor<32xi64, #sparse{{[0-9]*}}>,
290// CHECK-SAME:             %[[VAL_1:.*]]: tensor<32xi64>,
291// CHECK-SAME:             %[[VAL_2:.*]]: tensor<32xi64>) -> tensor<32xi64> {
292// CHECK-DAG:           %[[VAL_3:.*]] = arith.constant 32 : index
293// CHECK-DAG:           %[[VAL_4:.*]] = arith.constant 0 : index
294// CHECK-DAG:           %[[VAL_5:.*]] = arith.constant true
295// CHECK-DAG:           %[[VAL_6:.*]] = arith.constant 1 : index
296// CHECK-DAG:           %[[VAL_7:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}> to memref<?xindex>
297// CHECK-DAG:           %[[VAL_8:.*]] = sparse_tensor.coordinates %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}> to memref<?xindex>
298// CHECK-DAG:           %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xi64, #sparse{{[0-9]*}}> to memref<?xi64>
299// CHECK-DAG:           %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_1]] : tensor<32xi64> to memref<32xi64>
300// CHECK-DAG:           %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : tensor<32xi64> to memref<32xi64>
301// CHECK:           %[[VAL_12:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_4]]] : memref<?xindex>
302// CHECK:           %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex>
303// CHECK:           %[[VAL_14:.*]]:2 = scf.while (%[[VAL_15:.*]] = %[[VAL_12]], %[[VAL_16:.*]] = %[[VAL_4]]) : (index, index) -> (index, index) {
304// CHECK:             %[[VAL_17:.*]] = arith.cmpi ult, %[[VAL_15]], %[[VAL_13]] : index
305// CHECK:             scf.condition(%[[VAL_17]]) %[[VAL_15]], %[[VAL_16]] : index, index
306// CHECK:           } do {
307// CHECK:           ^bb0(%[[VAL_18:.*]]: index, %[[VAL_19:.*]]: index):
308// CHECK:             %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_18]]] : memref<?xindex>
309// CHECK:             %[[VAL_21:.*]] = arith.cmpi eq, %[[VAL_20]], %[[VAL_19]] : index
310// CHECK:             scf.if %[[VAL_21]] {
311// CHECK:               %[[VAL_22:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_18]]] : memref<?xi64>
312// CHECK:               %[[VAL_23:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_19]]] : memref<32xi64>
313// CHECK:               %[[VAL_24:.*]] = arith.ori %[[VAL_22]], %[[VAL_23]] : i64
314// CHECK:               memref.store %[[VAL_24]], %[[VAL_11]]{{\[}}%[[VAL_19]]] : memref<32xi64>
315// CHECK:             } else {
316// CHECK:               scf.if %[[VAL_5]] {
317// CHECK:                 %[[VAL_25:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_19]]] : memref<32xi64>
318// CHECK:                 memref.store %[[VAL_25]], %[[VAL_11]]{{\[}}%[[VAL_19]]] : memref<32xi64>
319// CHECK:               } else {
320// CHECK:               }
321// CHECK:             }
322// CHECK:             %[[VAL_26:.*]] = arith.cmpi eq, %[[VAL_20]], %[[VAL_19]] : index
323// CHECK:             %[[VAL_27:.*]] = arith.addi %[[VAL_18]], %[[VAL_6]] : index
324// CHECK:             %[[VAL_28:.*]] = arith.select %[[VAL_26]], %[[VAL_27]], %[[VAL_18]] : index
325// CHECK:             %[[VAL_29:.*]] = arith.addi %[[VAL_19]], %[[VAL_6]] : index
326// CHECK:             scf.yield %[[VAL_28]], %[[VAL_29]] : index, index
327// CHECK:           }
328// CHECK:           scf.for %[[VAL_30:.*]] = %[[VAL_31:.*]]#1 to %[[VAL_3]] step %[[VAL_6]] {
329// CHECK:             %[[VAL_32:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_30]]] : memref<32xi64>
330// CHECK:             memref.store %[[VAL_32]], %[[VAL_11]]{{\[}}%[[VAL_30]]] : memref<32xi64>
331// CHECK:           }
332// CHECK:           %[[VAL_33:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xi64>
333// CHECK:           return %[[VAL_33]] : tensor<32xi64>
334// CHECK:         }
335func.func @or(%arga: tensor<32xi64, #SV>,
336              %argb: tensor<32xi64>,
337              %argx: tensor<32xi64>) -> tensor<32xi64> {
338  %0 = linalg.generic #trait2
339     ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>)
340    outs(%argx: tensor<32xi64>) {
341      ^bb(%a: i64, %b: i64, %x: i64):
342        %0 = arith.ori %a, %b : i64
343        linalg.yield %0 : i64
344  } -> tensor<32xi64>
345  return %0 : tensor<32xi64>
346}
347
348// CHECK-LABEL:   func @xor(
349// CHECK-SAME:             %[[VAL_0:.*]]: tensor<32xi64, #sparse{{[0-9]*}}>,
350// CHECK-SAME:             %[[VAL_1:.*]]: tensor<32xi64>,
351// CHECK-SAME:             %[[VAL_2:.*]]: tensor<32xi64>) -> tensor<32xi64> {
352// CHECK-DAG:           %[[VAL_3:.*]] = arith.constant 32 : index
353// CHECK-DAG:           %[[VAL_4:.*]] = arith.constant 0 : index
354// CHECK-DAG:           %[[VAL_5:.*]] = arith.constant true
355// CHECK-DAG:           %[[VAL_6:.*]] = arith.constant 1 : index
356// CHECK-DAG:           %[[VAL_7:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}> to memref<?xindex>
357// CHECK-DAG:           %[[VAL_8:.*]] = sparse_tensor.coordinates %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}> to memref<?xindex>
358// CHECK-DAG:           %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xi64, #sparse{{[0-9]*}}> to memref<?xi64>
359// CHECK-DAG:           %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_1]] : tensor<32xi64> to memref<32xi64>
360// CHECK-DAG:           %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : tensor<32xi64> to memref<32xi64>
361// CHECK:           %[[VAL_12:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_4]]] : memref<?xindex>
362// CHECK:           %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex>
363// CHECK:           %[[VAL_14:.*]]:2 = scf.while (%[[VAL_15:.*]] = %[[VAL_12]], %[[VAL_16:.*]] = %[[VAL_4]]) : (index, index) -> (index, index) {
364// CHECK:             %[[VAL_17:.*]] = arith.cmpi ult, %[[VAL_15]], %[[VAL_13]] : index
365// CHECK:             scf.condition(%[[VAL_17]]) %[[VAL_15]], %[[VAL_16]] : index, index
366// CHECK:           } do {
367// CHECK:           ^bb0(%[[VAL_18:.*]]: index, %[[VAL_19:.*]]: index):
368// CHECK:             %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_18]]] : memref<?xindex>
369// CHECK:             %[[VAL_21:.*]] = arith.cmpi eq, %[[VAL_20]], %[[VAL_19]] : index
370// CHECK:             scf.if %[[VAL_21]] {
371// CHECK:               %[[VAL_22:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_18]]] : memref<?xi64>
372// CHECK:               %[[VAL_23:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_19]]] : memref<32xi64>
373// CHECK:               %[[VAL_24:.*]] = arith.xori %[[VAL_22]], %[[VAL_23]] : i64
374// CHECK:               memref.store %[[VAL_24]], %[[VAL_11]]{{\[}}%[[VAL_19]]] : memref<32xi64>
375// CHECK:             } else {
376// CHECK:               scf.if %[[VAL_5]] {
377// CHECK:                 %[[VAL_25:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_19]]] : memref<32xi64>
378// CHECK:                 memref.store %[[VAL_25]], %[[VAL_11]]{{\[}}%[[VAL_19]]] : memref<32xi64>
379// CHECK:               } else {
380// CHECK:               }
381// CHECK:             }
382// CHECK:             %[[VAL_26:.*]] = arith.cmpi eq, %[[VAL_20]], %[[VAL_19]] : index
383// CHECK:             %[[VAL_27:.*]] = arith.addi %[[VAL_18]], %[[VAL_6]] : index
384// CHECK:             %[[VAL_28:.*]] = arith.select %[[VAL_26]], %[[VAL_27]], %[[VAL_18]] : index
385// CHECK:             %[[VAL_29:.*]] = arith.addi %[[VAL_19]], %[[VAL_6]] : index
386// CHECK:             scf.yield %[[VAL_28]], %[[VAL_29]] : index, index
387// CHECK:           }
388// CHECK:           scf.for %[[VAL_30:.*]] = %[[VAL_31:.*]]#1 to %[[VAL_3]] step %[[VAL_6]] {
389// CHECK:             %[[VAL_32:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_30]]] : memref<32xi64>
390// CHECK:             memref.store %[[VAL_32]], %[[VAL_11]]{{\[}}%[[VAL_30]]] : memref<32xi64>
391// CHECK:           }
392// CHECK:           %[[VAL_33:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xi64>
393// CHECK:           return %[[VAL_33]] : tensor<32xi64>
394// CHECK:         }
395func.func @xor(%arga: tensor<32xi64, #SV>,
396               %argb: tensor<32xi64>,
397               %argx: tensor<32xi64>) -> tensor<32xi64> {
398  %0 = linalg.generic #trait2
399     ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>)
400    outs(%argx: tensor<32xi64>) {
401      ^bb(%a: i64, %b: i64, %x: i64):
402        %0 = arith.xori %a, %b : i64
403        linalg.yield %0 : i64
404  } -> tensor<32xi64>
405  return %0 : tensor<32xi64>
406}
407
408// CHECK-LABEL:   func @ashrbyc(
409// CHECK-SAME:                  %[[VAL_0:.*]]: tensor<32xi64, #sparse{{[0-9]*}}>,
410// CHECK-SAME:                  %[[VAL_1:.*]]: tensor<32xi64>) -> tensor<32xi64> {
411// CHECK-DAG:           %[[VAL_2:.*]] = arith.constant 2 : i64
412// CHECK-DAG:           %[[VAL_3:.*]] = arith.constant 0 : index
413// CHECK-DAG:           %[[VAL_4:.*]] = arith.constant 1 : index
414// CHECK-DAG:           %[[VAL_5:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}> to memref<?xindex>
415// CHECK-DAG:           %[[VAL_6:.*]] = sparse_tensor.coordinates %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}> to memref<?xindex>
416// CHECK-DAG:           %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xi64, #sparse{{[0-9]*}}> to memref<?xi64>
417// CHECK-DAG:           %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]] : tensor<32xi64> to memref<32xi64>
418// CHECK:           %[[VAL_9:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
419// CHECK:           %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
420// CHECK:           scf.for %[[VAL_11:.*]] = %[[VAL_9]] to %[[VAL_10]] step %[[VAL_4]] {
421// CHECK:             %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xindex>
422// CHECK:             %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_11]]] : memref<?xi64>
423// CHECK:             %[[VAL_14:.*]] = arith.shrsi %[[VAL_13]], %[[VAL_2]] : i64
424// CHECK:             memref.store %[[VAL_14]], %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<32xi64>
425// CHECK:           }
426// CHECK:           %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xi64>
427// CHECK:           return %[[VAL_15]] : tensor<32xi64>
428// CHECK:         }
429func.func @ashrbyc(%arga: tensor<32xi64, #SV>,
430                   %argx: tensor<32xi64>) -> tensor<32xi64> {
431  %c = arith.constant 2 : i64
432  %0 = linalg.generic #traitc
433     ins(%arga: tensor<32xi64, #SV>)
434    outs(%argx: tensor<32xi64>) {
435      ^bb(%a: i64, %x: i64):
436        %0 = arith.shrsi %a, %c : i64
437        linalg.yield %0 : i64
438  } -> tensor<32xi64>
439  return %0 : tensor<32xi64>
440}
441
442// CHECK-LABEL:   func @lsrbyc(
443// CHECK-SAME:                 %[[VAL_0:.*]]: tensor<32xi64, #sparse{{[0-9]*}}>,
444// CHECK-SAME:                 %[[VAL_1:.*]]: tensor<32xi64>) -> tensor<32xi64> {
445// CHECK-DAG:           %[[VAL_2:.*]] = arith.constant 2 : i64
446// CHECK-DAG:           %[[VAL_3:.*]] = arith.constant 0 : index
447// CHECK-DAG:           %[[VAL_4:.*]] = arith.constant 1 : index
448// CHECK-DAG:           %[[VAL_5:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}> to memref<?xindex>
449// CHECK-DAG:           %[[VAL_6:.*]] = sparse_tensor.coordinates %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}> to memref<?xindex>
450// CHECK-DAG:           %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xi64, #sparse{{[0-9]*}}> to memref<?xi64>
451// CHECK-DAG:           %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]] : tensor<32xi64> to memref<32xi64>
452// CHECK:           %[[VAL_9:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
453// CHECK:           %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
454// CHECK:           scf.for %[[VAL_11:.*]] = %[[VAL_9]] to %[[VAL_10]] step %[[VAL_4]] {
455// CHECK:             %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xindex>
456// CHECK:             %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_11]]] : memref<?xi64>
457// CHECK:             %[[VAL_14:.*]] = arith.shrui %[[VAL_13]], %[[VAL_2]] : i64
458// CHECK:             memref.store %[[VAL_14]], %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<32xi64>
459// CHECK:           }
460// CHECK:           %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xi64>
461// CHECK:           return %[[VAL_15]] : tensor<32xi64>
462// CHECK:         }
463func.func @lsrbyc(%arga: tensor<32xi64, #SV>,
464                  %argx: tensor<32xi64>) -> tensor<32xi64> {
465  %c = arith.constant 2 : i64
466  %0 = linalg.generic #traitc
467     ins(%arga: tensor<32xi64, #SV>)
468    outs(%argx: tensor<32xi64>) {
469      ^bb(%a: i64, %x: i64):
470        %0 = arith.shrui %a, %c : i64
471        linalg.yield %0 : i64
472  } -> tensor<32xi64>
473  return %0 : tensor<32xi64>
474}
475
476// CHECK-LABEL:   func @lslbyc(
477// CHECK-SAME:                 %[[VAL_0:.*]]: tensor<32xi64, #sparse{{[0-9]*}}>,
478// CHECK-SAME:                 %[[VAL_1:.*]]: tensor<32xi64>) -> tensor<32xi64> {
479// CHECK-DAG:           %[[VAL_2:.*]] = arith.constant 2 : i64
480// CHECK-DAG:           %[[VAL_3:.*]] = arith.constant 0 : index
481// CHECK-DAG:           %[[VAL_4:.*]] = arith.constant 1 : index
482// CHECK-DAG:           %[[VAL_5:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}> to memref<?xindex>
483// CHECK-DAG:           %[[VAL_6:.*]] = sparse_tensor.coordinates %[[VAL_0]] {level = 0 : index} : tensor<32xi64, #sparse{{[0-9]*}}> to memref<?xindex>
484// CHECK-DAG:           %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xi64, #sparse{{[0-9]*}}> to memref<?xi64>
485// CHECK-DAG:           %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]] : tensor<32xi64> to memref<32xi64>
486// CHECK:           %[[VAL_9:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
487// CHECK:           %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
488// CHECK:           scf.for %[[VAL_11:.*]] = %[[VAL_9]] to %[[VAL_10]] step %[[VAL_4]] {
489// CHECK:             %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xindex>
490// CHECK:             %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_11]]] : memref<?xi64>
491// CHECK:             %[[VAL_14:.*]] = arith.shli %[[VAL_13]], %[[VAL_2]] : i64
492// CHECK:             memref.store %[[VAL_14]], %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<32xi64>
493// CHECK:           }
494// CHECK:           %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xi64>
495// CHECK:           return %[[VAL_15]] : tensor<32xi64>
496// CHECK:         }
497func.func @lslbyc(%arga: tensor<32xi64, #SV>,
498                  %argx: tensor<32xi64>) -> tensor<32xi64> {
499  %c = arith.constant 2 : i64
500  %0 = linalg.generic #traitc
501     ins(%arga: tensor<32xi64, #SV>)
502    outs(%argx: tensor<32xi64>) {
503      ^bb(%a: i64, %x: i64):
504        %0 = arith.shli %a, %c : i64
505        linalg.yield %0 : i64
506  } -> tensor<32xi64>
507  return %0 : tensor<32xi64>
508}
509