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
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
|
// RUN: mlir-opt %s -test-affine-reify-value-bounds="reify-to-func-args" \
// RUN: -verify-diagnostics -split-input-file | FileCheck %s
// RUN: mlir-opt %s -test-affine-reify-value-bounds="reify-to-func-args use-arith-ops" \
// RUN: -verify-diagnostics -split-input-file | FileCheck %s --check-prefix=CHECK-ARITH
// CHECK-LABEL: func @reify_through_chain(
// CHECK-SAME: %[[sz0:.*]]: index, %[[sz2:.*]]: index
// CHECK: %[[c10:.*]] = arith.constant 10 : index
// CHECK: return %[[sz0]], %[[c10]], %[[sz2]]
// CHECK-ARITH-LABEL: func @reify_through_chain(
// CHECK-ARITH-SAME: %[[sz0:.*]]: index, %[[sz2:.*]]: index
// CHECK-ARITH: %[[c10:.*]] = arith.constant 10 : index
// CHECK-ARITH: return %[[sz0]], %[[c10]], %[[sz2]]
func.func @reify_through_chain(%sz0: index, %sz2: index) -> (index, index, index) {
%c2 = arith.constant 2 : index
%0 = tensor.empty(%sz0, %sz2) : tensor<?x10x?xf32>
%1 = tensor.cast %0 : tensor<?x10x?xf32> to tensor<?x?x?xf32>
%pos = arith.constant 0 : index
%f = arith.constant 0.0 : f32
%2 = tensor.insert %f into %1[%pos, %pos, %pos] : tensor<?x?x?xf32>
%3 = tensor.dim %2, %c2 : tensor<?x?x?xf32>
%4 = "test.reify_bound"(%2) {dim = 0} : (tensor<?x?x?xf32>) -> (index)
%5 = "test.reify_bound"(%2) {dim = 1} : (tensor<?x?x?xf32>) -> (index)
%6 = "test.reify_bound"(%3) : (index) -> (index)
return %4, %5, %6 : index, index, index
}
// -----
// CHECK-LABEL: func @reify_slice_bound(
// CHECK: %[[c5:.*]] = arith.constant 5 : index
// CHECK: "test.some_use"(%[[c5]])
// CHECK: %[[c5:.*]] = arith.constant 5 : index
// CHECK: "test.some_use"(%[[c5]])
func.func @reify_slice_bound(%t: tensor<?x?xi32>, %idx: index, %ub: index, %f: f32) {
%c0 = arith.constant 0 : index
%c4 = arith.constant 4 : index
scf.for %iv = %c0 to %ub step %c4 {
%sz = affine.min affine_map<(d0)[s0] -> (-d0 + s0, 4)>(%iv)[%ub]
%slice = tensor.extract_slice %t[%idx, %iv] [1, %sz] [1, 1] : tensor<?x?xi32> to tensor<1x?xi32>
%filled = linalg.fill ins(%f : f32) outs(%slice : tensor<1x?xi32>) -> tensor<1x?xi32>
%bound = "test.reify_bound"(%filled) {dim = 1, type = "UB"} : (tensor<1x?xi32>) -> (index)
"test.some_use"(%bound) : (index) -> ()
%bound_const = "test.reify_constant_bound"(%filled) {dim = 1, type = "UB"} : (tensor<1x?xi32>) -> (index)
"test.some_use"(%bound_const) : (index) -> ()
}
return
}
// -----
// CHECK: #[[$map:.*]] = affine_map<()[s0, s1] -> (s0 - s1 + 1)>
// CHECK-LABEL: func @scf_for(
// CHECK-SAME: %[[lb:.*]]: index, %[[ub:.*]]: index, %[[step:.*]]: index
// CHECK: %[[bound:.*]] = affine.apply #[[$map]]()[%[[ub]], %[[lb]]]
// CHECK: "test.some_use"(%[[bound]])
func.func @scf_for(%lb: index, %ub: index, %step: index) {
scf.for %iv = %lb to %ub step %step {
%0 = affine.apply affine_map<(d0)[s0] -> (-d0 + s0)>(%iv)[%ub]
%bound = "test.reify_bound"(%0) {type = "UB"} : (index) -> (index)
"test.some_use"(%bound) : (index) -> ()
}
return
}
// -----
// CHECK-LABEL: func @reify_slice_bound2(
func.func @reify_slice_bound2(%lb0: index, %ub0: index, %step0: index,
%ub2: index, %t1: tensor<1x?xi8>,
%t2: tensor<?x?xi8>, %t3: tensor<1x?xi32>) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c32 = arith.constant 32 : index
scf.for %iv0 = %lb0 to %ub0 step %step0 {
// CHECK: %[[c129:.*]] = arith.constant 129 : index
// CHECK: "test.some_use"(%[[c129]])
%ub1 = affine.min affine_map<(d0)[s0] -> (-d0 + s0, 128)>(%iv0)[%ub0]
%ub1_ub = "test.reify_bound"(%ub1) {type = "UB"} : (index) -> (index)
"test.some_use"(%ub1_ub) : (index) -> ()
// CHECK: %[[c129:.*]] = arith.constant 129 : index
// CHECK: "test.some_use"(%[[c129]])
%lb1 = affine.apply affine_map<()[s0] -> ((s0 floordiv 32) * 32)>()[%ub1]
%lb1_ub = "test.reify_bound"(%lb1) {type = "UB"} : (index) -> (index)
"test.some_use"(%lb1_ub) : (index) -> ()
// CHECK: %[[c129:.*]] = arith.constant 129 : index
// CHECK: "test.some_use"(%[[c129]])
%lb1_ub_const = "test.reify_constant_bound"(%lb1) {type = "UB"} : (index) -> (index)
"test.some_use"(%lb1_ub_const) : (index) -> ()
scf.for %iv1 = %lb1 to %ub1 step %c32 {
// CHECK: %[[c32:.*]] = arith.constant 32 : index
// CHECK: "test.some_use"(%[[c32]])
%sz = affine.apply affine_map<(d0)[s0] -> (-d0 + s0)>(%iv1)[%ub1]
%sz_ub = "test.reify_bound"(%sz) {type = "UB"} : (index) -> (index)
"test.some_use"(%sz_ub) : (index) -> ()
scf.for %iv2 = %c0 to %ub2 step %c1 {
%slice1 = tensor.extract_slice %t1[0, %iv2] [1, 1] [1, 1] : tensor<1x?xi8> to tensor<1x1xi8>
%slice2 = tensor.extract_slice %t2[%iv2, 0] [1, %sz] [1, 1] : tensor<?x?xi8> to tensor<1x?xi8>
%slice3 = tensor.extract_slice %t3[0, 0] [1, %sz] [1, 1] : tensor<1x?xi32> to tensor<1x?xi32>
%matmul = linalg.matmul ins(%slice1, %slice2 : tensor<1x1xi8>, tensor<1x?xi8>) outs(%slice3 : tensor<1x?xi32>) -> tensor<1x?xi32>
// CHECK: %[[c32:.*]] = arith.constant 32 : index
// CHECK: "test.some_use"(%[[c32]])
%matmul_ub = "test.reify_bound"(%matmul) {dim = 1, type = "UB"} : (tensor<1x?xi32>) -> (index)
"test.some_use"(%matmul_ub) : (index) -> ()
// CHECK: %[[c32:.*]] = arith.constant 32 : index
// CHECK: "test.some_use"(%[[c32]])
%matmul_ub_const = "test.reify_constant_bound"(%matmul) {dim = 1, type = "UB"} : (tensor<1x?xi32>) -> (index)
"test.some_use"(%matmul_ub_const) : (index) -> ()
}
}
}
return
}
|
