blob: c68c576766b8566aef640edd6cf912d22cb61c17 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
|
// RUN: mlir-opt %s -sparsification | FileCheck %s
// The file contains examples that will be rejected by sparsifier
// (we expect the linalg.generic unchanged).
#SparseVector = #sparse_tensor.encoding<{map = (d0) -> (d0 : compressed)}>
#trait = {
indexing_maps = [
affine_map<(i) -> (i)>, // a (in)
affine_map<(i) -> ()> // x (out)
],
iterator_types = ["reduction"]
}
// CHECK-LABEL: func.func @sparse_reduction_subi(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<i32>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<?xi32, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<i32> {
// CHECK: %[[VAL_2:.*]] = linalg.generic
// CHECK: ^bb0(%[[VAL_3:.*]]: i32, %[[VAL_4:.*]]: i32):
// CHECK: %[[VAL_5:.*]] = arith.subi %[[VAL_3]], %[[VAL_4]] : i32
// CHECK: linalg.yield %[[VAL_5]] : i32
// CHECK: } -> tensor<i32>
// CHECK: return %[[VAL_6:.*]] : tensor<i32>
func.func @sparse_reduction_subi(%argx: tensor<i32>,
%arga: tensor<?xi32, #SparseVector>)
-> tensor<i32> {
%0 = linalg.generic #trait
ins(%arga: tensor<?xi32, #SparseVector>)
outs(%argx: tensor<i32>) {
^bb(%a: i32, %x: i32):
// NOTE: `subi %a, %x` is the reason why the program is rejected by the sparsifier.
// It is because we do not allow `-outTensor` in reduction loops as it creates cyclic
// dependences.
%t = arith.subi %a, %x: i32
linalg.yield %t : i32
} -> tensor<i32>
return %0 : tensor<i32>
}
|
