xref: /llvm-project/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir (revision 0aa831e0edb1c1deabb96ce2435667cc82bac79b)
1// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(test-affine-reify-value-bounds))' -verify-diagnostics \
2// RUN:     -verify-diagnostics -split-input-file | FileCheck %s
3
4// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(test-affine-reify-value-bounds{use-arith-ops}))' \
5// RUN:     -verify-diagnostics -split-input-file | \
6// RUN: FileCheck %s --check-prefix=CHECK-ARITH
7
8// CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 + 5)>
9// CHECK-LABEL: func @arith_addi(
10//  CHECK-SAME:     %[[a:.*]]: index
11//       CHECK:   %[[apply:.*]] = affine.apply #[[$map]]()[%[[a]]]
12//       CHECK:   return %[[apply]]
13
14// CHECK-ARITH-LABEL: func @arith_addi(
15//  CHECK-ARITH-SAME:     %[[a:.*]]: index
16//       CHECK-ARITH:   %[[c5:.*]] = arith.constant 5 : index
17//       CHECK-ARITH:   %[[add:.*]] = arith.addi %[[c5]], %[[a]]
18//       CHECK-ARITH:   %[[c5:.*]] = arith.constant 5 : index
19//       CHECK-ARITH:   %[[add:.*]] = arith.addi %[[a]], %[[c5]]
20//       CHECK-ARITH:   return %[[add]]
21func.func @arith_addi(%a: index) -> index {
22  %0 = arith.constant 5 : index
23  %1 = arith.addi %0, %a : index
24  %2 = "test.reify_bound"(%1) : (index) -> (index)
25  return %2 : index
26}
27
28// -----
29
30// CHECK: #[[$map:.*]] = affine_map<()[s0] -> (-s0 + 5)>
31// CHECK-LABEL: func @arith_subi(
32//  CHECK-SAME:     %[[a:.*]]: index
33//       CHECK:   %[[apply:.*]] = affine.apply #[[$map]]()[%[[a]]]
34//       CHECK:   return %[[apply]]
35func.func @arith_subi(%a: index) -> index {
36  %0 = arith.constant 5 : index
37  %1 = arith.subi %0, %a : index
38  %2 = "test.reify_bound"(%1) : (index) -> (index)
39  return %2 : index
40}
41
42// -----
43
44// CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 * 5)>
45// CHECK-LABEL: func @arith_muli(
46//  CHECK-SAME:     %[[a:.*]]: index
47//       CHECK:   %[[apply:.*]] = affine.apply #[[$map]]()[%[[a]]]
48//       CHECK:   return %[[apply]]
49func.func @arith_muli(%a: index) -> index {
50  %0 = arith.constant 5 : index
51  %1 = arith.muli %0, %a : index
52  %2 = "test.reify_bound"(%1) : (index) -> (index)
53  return %2 : index
54}
55
56// -----
57
58func.func @arith_muli_non_pure(%a: index, %b: index) -> index {
59  %0 = arith.muli %a, %b : index
60  // Semi-affine expressions (such as "symbol * symbol") are not supported.
61  // expected-error @below{{could not reify bound}}
62  %1 = "test.reify_bound"(%0) : (index) -> (index)
63  return %1 : index
64}
65
66// -----
67
68// CHECK-LABEL: func @arith_const()
69//       CHECK:   %[[c5:.*]] = arith.constant 5 : index
70//       CHECK:   %[[c5:.*]] = arith.constant 5 : index
71//       CHECK:   return %[[c5]]
72func.func @arith_const() -> index {
73  %c5 = arith.constant 5 : index
74  %0 = "test.reify_bound"(%c5) : (index) -> (index)
75  return %0 : index
76}
77
78// -----
79
80// CHECK-LABEL: func @arith_select(
81func.func @arith_select(%c: i1) -> (index, index) {
82  // CHECK: arith.constant 5 : index
83  %c5 = arith.constant 5 : index
84  // CHECK: arith.constant 9 : index
85  %c9 = arith.constant 9 : index
86  %r = arith.select %c, %c5, %c9 : index
87  // CHECK: %[[c5:.*]] = arith.constant 5 : index
88  // CHECK: %[[c10:.*]] = arith.constant 10 : index
89  %0 = "test.reify_bound"(%r) {type = "LB"} : (index) -> (index)
90  %1 = "test.reify_bound"(%r) {type = "UB"} : (index) -> (index)
91  // CHECK: return %[[c5]], %[[c10]]
92  return %0, %1 : index, index
93}
94
95// -----
96
97// CHECK-LABEL: func @arith_select_elementwise(
98//  CHECK-SAME:     %[[a:.*]]: tensor<?xf32>, %[[b:.*]]: tensor<?xf32>, %[[c:.*]]: tensor<?xi1>)
99func.func @arith_select_elementwise(%a: tensor<?xf32>, %b: tensor<?xf32>, %c: tensor<?xi1>) -> index {
100  %r = arith.select %c, %a, %b : tensor<?xi1>, tensor<?xf32>
101  // CHECK: %[[c0:.*]] = arith.constant 0 : index
102  // CHECK: %[[dim:.*]] = tensor.dim %[[a]], %[[c0]]
103  %0 = "test.reify_bound"(%r) {type = "EQ", dim = 0}
104      : (tensor<?xf32>) -> (index)
105  // CHECK: return %[[dim]]
106  return %0 : index
107}
108