xref: /llvm-project/mlir/test/Dialect/Math/algebraic-simplification.mlir (revision bb6d5c220004a5d7e466a669324001285a688918)
1// RUN: mlir-opt %s -test-math-algebraic-simplification | FileCheck %s --dump-input=always
2
3// CHECK-LABEL: @pow_noop
4func.func @pow_noop(%arg0: f32, %arg1 : vector<4xf32>) -> (f32, vector<4xf32>) {
5  // CHECK: return %arg0, %arg1
6  %c = arith.constant 1.0 : f32
7  %v = arith.constant dense <1.0> : vector<4xf32>
8  %0 = math.powf %arg0, %c : f32
9  %1 = math.powf %arg1, %v : vector<4xf32>
10  return %0, %1 : f32, vector<4xf32>
11}
12
13// CHECK-LABEL: @pow_square
14func.func @pow_square(%arg0: f32, %arg1 : vector<4xf32>) -> (f32, vector<4xf32>) {
15  // CHECK: %[[SCALAR:.*]] = arith.mulf %arg0, %arg0
16  // CHECK: %[[VECTOR:.*]] = arith.mulf %arg1, %arg1
17  // CHECK: return %[[SCALAR]], %[[VECTOR]]
18  %c = arith.constant 2.0 : f32
19  %v = arith.constant dense <2.0> : vector<4xf32>
20  %0 = math.powf %arg0, %c : f32
21  %1 = math.powf %arg1, %v : vector<4xf32>
22  return %0, %1 : f32, vector<4xf32>
23}
24
25// CHECK-LABEL: @pow_cube
26func.func @pow_cube(%arg0: f32, %arg1 : vector<4xf32>) -> (f32, vector<4xf32>) {
27  // CHECK: %[[TMP_S:.*]] = arith.mulf %arg0, %arg0
28  // CHECK: %[[SCALAR:.*]] = arith.mulf %arg0, %[[TMP_S]]
29  // CHECK: %[[TMP_V:.*]] = arith.mulf %arg1, %arg1
30  // CHECK: %[[VECTOR:.*]] = arith.mulf %arg1, %[[TMP_V]]
31  // CHECK: return %[[SCALAR]], %[[VECTOR]]
32  %c = arith.constant 3.0 : f32
33  %v = arith.constant dense <3.0> : vector<4xf32>
34  %0 = math.powf %arg0, %c : f32
35  %1 = math.powf %arg1, %v : vector<4xf32>
36  return %0, %1 : f32, vector<4xf32>
37}
38
39// CHECK-LABEL: @pow_recip
40func.func @pow_recip(%arg0: f32, %arg1 : vector<4xf32>) -> (f32, vector<4xf32>) {
41  // CHECK-DAG: %[[CST_S:.*]] = arith.constant 1.0{{.*}} : f32
42  // CHECK-DAG: %[[CST_V:.*]] = arith.constant dense<1.0{{.*}}> : vector<4xf32>
43  // CHECK: %[[SCALAR:.*]] = arith.divf %[[CST_S]], %arg0
44  // CHECK: %[[VECTOR:.*]] = arith.divf %[[CST_V]], %arg1
45  // CHECK: return %[[SCALAR]], %[[VECTOR]]
46  %c = arith.constant -1.0 : f32
47  %v = arith.constant dense <-1.0> : vector<4xf32>
48  %0 = math.powf %arg0, %c : f32
49  %1 = math.powf %arg1, %v : vector<4xf32>
50  return %0, %1 : f32, vector<4xf32>
51}
52
53// CHECK-LABEL: @pow_sqrt
54func.func @pow_sqrt(%arg0: f32, %arg1 : vector<4xf32>) -> (f32, vector<4xf32>) {
55  // CHECK: %[[SCALAR:.*]] = math.sqrt %arg0
56  // CHECK: %[[VECTOR:.*]] = math.sqrt %arg1
57  // CHECK: return %[[SCALAR]], %[[VECTOR]]
58  %c = arith.constant 0.5 : f32
59  %v = arith.constant dense <0.5> : vector<4xf32>
60  %0 = math.powf %arg0, %c : f32
61  %1 = math.powf %arg1, %v : vector<4xf32>
62  return %0, %1 : f32, vector<4xf32>
63}
64
65// CHECK-LABEL: @pow_rsqrt
66func.func @pow_rsqrt(%arg0: f32, %arg1 : vector<4xf32>) -> (f32, vector<4xf32>) {
67  // CHECK: %[[SCALAR:.*]] = math.rsqrt %arg0
68  // CHECK: %[[VECTOR:.*]] = math.rsqrt %arg1
69  // CHECK: return %[[SCALAR]], %[[VECTOR]]
70  %c = arith.constant -0.5 : f32
71  %v = arith.constant dense <-0.5> : vector<4xf32>
72  %0 = math.powf %arg0, %c : f32
73  %1 = math.powf %arg1, %v : vector<4xf32>
74  return %0, %1 : f32, vector<4xf32>
75}
76
77// CHECK-LABEL: @pow_0_75
78func.func @pow_0_75(%arg0: f32, %arg1 : vector<4xf32>) -> (f32, vector<4xf32>) {
79  // CHECK: %[[SQRT1S:.*]] = math.sqrt %arg0
80  // CHECK: %[[SQRT2S:.*]] = math.sqrt %[[SQRT1S]]
81  // CHECK: %[[SCALAR:.*]] = arith.mulf %[[SQRT1S]], %[[SQRT2S]]
82  // CHECK: %[[SQRT1V:.*]] = math.sqrt %arg1
83  // CHECK: %[[SQRT2V:.*]] = math.sqrt %[[SQRT1V]]
84  // CHECK: %[[VECTOR:.*]] = arith.mulf %[[SQRT1V]], %[[SQRT2V]]
85  // CHECK: return %[[SCALAR]], %[[VECTOR]]
86  %c = arith.constant 0.75 : f32
87  %v = arith.constant dense <0.75> : vector<4xf32>
88  %0 = math.powf %arg0, %c : f32
89  %1 = math.powf %arg1, %v : vector<4xf32>
90  return %0, %1 : f32, vector<4xf32>
91}
92
93// CHECK-LABEL: @ipowi_zero_exp(
94// CHECK-SAME: %[[ARG0:.+]]: i32
95// CHECK-SAME: %[[ARG1:.+]]: vector<4xi32>
96// CHECK-SAME: -> (i32, vector<4xi32>) {
97func.func @ipowi_zero_exp(%arg0: i32, %arg1: vector<4xi32>) -> (i32, vector<4xi32>) {
98  // CHECK: %[[CST_S:.*]] = arith.constant 1 : i32
99  // CHECK: %[[CST_V:.*]] = arith.constant dense<1> : vector<4xi32>
100  // CHECK: return %[[CST_S]], %[[CST_V]]
101  %c = arith.constant 0 : i32
102  %v = arith.constant dense <0> : vector<4xi32>
103  %0 = math.ipowi %arg0, %c : i32
104  %1 = math.ipowi %arg1, %v : vector<4xi32>
105  return %0, %1 : i32, vector<4xi32>
106}
107
108// CHECK-LABEL: @ipowi_exp_one(
109// CHECK-SAME: %[[ARG0:.+]]: i32
110// CHECK-SAME: %[[ARG1:.+]]: vector<4xi32>
111// CHECK-SAME: -> (i32, vector<4xi32>, i32, vector<4xi32>) {
112func.func @ipowi_exp_one(%arg0: i32, %arg1: vector<4xi32>) -> (i32, vector<4xi32>, i32, vector<4xi32>) {
113  // CHECK-DAG: %[[CST_S:.*]] = arith.constant 1 : i32
114  // CHECK-DAG: %[[CST_V:.*]] = arith.constant dense<1> : vector<4xi32>
115  // CHECK: %[[SCALAR:.*]] = arith.divsi %[[CST_S]], %[[ARG0]]
116  // CHECK: %[[VECTOR:.*]] = arith.divsi %[[CST_V]], %[[ARG1]]
117  // CHECK: return %[[ARG0]], %[[ARG1]], %[[SCALAR]], %[[VECTOR]]
118  %c1 = arith.constant 1 : i32
119  %v1 = arith.constant dense <1> : vector<4xi32>
120  %0 = math.ipowi %arg0, %c1 : i32
121  %1 = math.ipowi %arg1, %v1 : vector<4xi32>
122  %cm1 = arith.constant -1 : i32
123  %vm1 = arith.constant dense <-1> : vector<4xi32>
124  %2 = math.ipowi %arg0, %cm1 : i32
125  %3 = math.ipowi %arg1, %vm1 : vector<4xi32>
126  return %0, %1, %2, %3 : i32, vector<4xi32>, i32, vector<4xi32>
127}
128
129// CHECK-LABEL: @ipowi_exp_two(
130// CHECK-SAME: %[[ARG0:.+]]: i32
131// CHECK-SAME: %[[ARG1:.+]]: vector<4xi32>
132// CHECK-SAME: -> (i32, vector<4xi32>, i32, vector<4xi32>) {
133func.func @ipowi_exp_two(%arg0: i32, %arg1: vector<4xi32>) -> (i32, vector<4xi32>, i32, vector<4xi32>) {
134  // CHECK-DAG: %[[CST_S:.*]] = arith.constant 1 : i32
135  // CHECK-DAG: %[[CST_V:.*]] = arith.constant dense<1> : vector<4xi32>
136  // CHECK: %[[SCALAR0:.*]] = arith.muli %[[ARG0]], %[[ARG0]]
137  // CHECK: %[[VECTOR0:.*]] = arith.muli %[[ARG1]], %[[ARG1]]
138  // CHECK: %[[SCALAR1:.*]] = arith.divsi %[[CST_S]], %[[ARG0]]
139  // CHECK: %[[SMUL:.*]] = arith.muli %[[SCALAR1]], %[[SCALAR1]]
140  // CHECK: %[[VECTOR1:.*]] = arith.divsi %[[CST_V]], %[[ARG1]]
141  // CHECK: %[[VMUL:.*]] = arith.muli %[[VECTOR1]], %[[VECTOR1]]
142  // CHECK: return %[[SCALAR0]], %[[VECTOR0]], %[[SMUL]], %[[VMUL]]
143  %c1 = arith.constant 2 : i32
144  %v1 = arith.constant dense <2> : vector<4xi32>
145  %0 = math.ipowi %arg0, %c1 : i32
146  %1 = math.ipowi %arg1, %v1 : vector<4xi32>
147  %cm1 = arith.constant -2 : i32
148  %vm1 = arith.constant dense <-2> : vector<4xi32>
149  %2 = math.ipowi %arg0, %cm1 : i32
150  %3 = math.ipowi %arg1, %vm1 : vector<4xi32>
151  return %0, %1, %2, %3 : i32, vector<4xi32>, i32, vector<4xi32>
152}
153
154// CHECK-LABEL: @ipowi_exp_three(
155// CHECK-SAME: %[[ARG0:.+]]: i32
156// CHECK-SAME: %[[ARG1:.+]]: vector<4xi32>
157// CHECK-SAME: -> (i32, vector<4xi32>, i32, vector<4xi32>) {
158func.func @ipowi_exp_three(%arg0: i32, %arg1: vector<4xi32>) -> (i32, vector<4xi32>, i32, vector<4xi32>) {
159  // CHECK-DAG: %[[CST_S:.*]] = arith.constant 1 : i32
160  // CHECK-DAG: %[[CST_V:.*]] = arith.constant dense<1> : vector<4xi32>
161  // CHECK: %[[SMUL0:.*]] = arith.muli %[[ARG0]], %[[ARG0]]
162  // CHECK: %[[SCALAR0:.*]] = arith.muli %[[SMUL0]], %[[ARG0]]
163  // CHECK: %[[VMUL0:.*]] = arith.muli %[[ARG1]], %[[ARG1]]
164  // CHECK: %[[VECTOR0:.*]] = arith.muli %[[VMUL0]], %[[ARG1]]
165  // CHECK: %[[SCALAR1:.*]] = arith.divsi %[[CST_S]], %[[ARG0]]
166  // CHECK: %[[SMUL1:.*]] = arith.muli %[[SCALAR1]], %[[SCALAR1]]
167  // CHECK: %[[SMUL2:.*]] = arith.muli %[[SMUL1]], %[[SCALAR1]]
168  // CHECK: %[[VECTOR1:.*]] = arith.divsi %[[CST_V]], %[[ARG1]]
169  // CHECK: %[[VMUL1:.*]] = arith.muli %[[VECTOR1]], %[[VECTOR1]]
170  // CHECK: %[[VMUL2:.*]] = arith.muli %[[VMUL1]], %[[VECTOR1]]
171  // CHECK: return %[[SCALAR0]], %[[VECTOR0]], %[[SMUL2]], %[[VMUL2]]
172  %c1 = arith.constant 3 : i32
173  %v1 = arith.constant dense <3> : vector<4xi32>
174  %0 = math.ipowi %arg0, %c1 : i32
175  %1 = math.ipowi %arg1, %v1 : vector<4xi32>
176  %cm1 = arith.constant -3 : i32
177  %vm1 = arith.constant dense <-3> : vector<4xi32>
178  %2 = math.ipowi %arg0, %cm1 : i32
179  %3 = math.ipowi %arg1, %vm1 : vector<4xi32>
180  return %0, %1, %2, %3 : i32, vector<4xi32>, i32, vector<4xi32>
181}
182
183// CHECK-LABEL: @fpowi_zero_exp(
184// CHECK-SAME: %[[ARG0:.+]]: f32
185// CHECK-SAME: %[[ARG1:.+]]: vector<4xf32>
186// CHECK-SAME: -> (f32, vector<4xf32>) {
187func.func @fpowi_zero_exp(%arg0: f32, %arg1: vector<4xf32>) -> (f32, vector<4xf32>) {
188  // CHECK: %[[CST_S:.*]] = arith.constant 1.000000e+00 : f32
189  // CHECK: %[[CST_V:.*]] = arith.constant dense<1.000000e+00> : vector<4xf32>
190  // CHECK: return %[[CST_S]], %[[CST_V]]
191  %c = arith.constant 0 : i32
192  %v = arith.constant dense <0> : vector<4xi32>
193  %0 = math.fpowi %arg0, %c : f32, i32
194  %1 = math.fpowi %arg1, %v : vector<4xf32>, vector<4xi32>
195  return %0, %1 : f32, vector<4xf32>
196}
197
198// CHECK-LABEL: @fpowi_exp_one(
199// CHECK-SAME: %[[ARG0:.+]]: f32
200// CHECK-SAME: %[[ARG1:.+]]: vector<4xf32>
201// CHECK-SAME: -> (f32, vector<4xf32>, f32, vector<4xf32>) {
202func.func @fpowi_exp_one(%arg0: f32, %arg1: vector<4xf32>) -> (f32, vector<4xf32>, f32, vector<4xf32>) {
203  // CHECK-DAG: %[[CST_S:.*]] = arith.constant 1.000000e+00 : f32
204  // CHECK-DAG: %[[CST_V:.*]] = arith.constant dense<1.000000e+00> : vector<4xf32>
205  // CHECK: %[[SCALAR:.*]] = arith.divf %[[CST_S]], %[[ARG0]]
206  // CHECK: %[[VECTOR:.*]] = arith.divf %[[CST_V]], %[[ARG1]]
207  // CHECK: return %[[ARG0]], %[[ARG1]], %[[SCALAR]], %[[VECTOR]]
208  %c1 = arith.constant 1 : i32
209  %v1 = arith.constant dense <1> : vector<4xi32>
210  %0 = math.fpowi %arg0, %c1 : f32, i32
211  %1 = math.fpowi %arg1, %v1 : vector<4xf32>, vector<4xi32>
212  %cm1 = arith.constant -1 : i32
213  %vm1 = arith.constant dense <-1> : vector<4xi32>
214  %2 = math.fpowi %arg0, %cm1 : f32, i32
215  %3 = math.fpowi %arg1, %vm1 : vector<4xf32>, vector<4xi32>
216  return %0, %1, %2, %3 : f32, vector<4xf32>, f32, vector<4xf32>
217}
218
219// CHECK-LABEL: @fpowi_exp_two(
220// CHECK-SAME: %[[ARG0:.+]]: f32
221// CHECK-SAME: %[[ARG1:.+]]: vector<4xf32>
222// CHECK-SAME: -> (f32, vector<4xf32>, f32, vector<4xf32>) {
223func.func @fpowi_exp_two(%arg0: f32, %arg1: vector<4xf32>) -> (f32, vector<4xf32>, f32, vector<4xf32>) {
224  // CHECK-DAG: %[[CST_S:.*]] = arith.constant 1.000000e+00 : f32
225  // CHECK-DAG: %[[CST_V:.*]] = arith.constant dense<1.000000e+00> : vector<4xf32>
226  // CHECK: %[[SCALAR0:.*]] = arith.mulf %[[ARG0]], %[[ARG0]]
227  // CHECK: %[[VECTOR0:.*]] = arith.mulf %[[ARG1]], %[[ARG1]]
228  // CHECK: %[[SCALAR1:.*]] = arith.divf %[[CST_S]], %[[ARG0]]
229  // CHECK: %[[SMUL:.*]] = arith.mulf %[[SCALAR1]], %[[SCALAR1]]
230  // CHECK: %[[VECTOR1:.*]] = arith.divf %[[CST_V]], %[[ARG1]]
231  // CHECK: %[[VMUL:.*]] = arith.mulf %[[VECTOR1]], %[[VECTOR1]]
232  // CHECK: return %[[SCALAR0]], %[[VECTOR0]], %[[SMUL]], %[[VMUL]]
233  %c1 = arith.constant 2 : i32
234  %v1 = arith.constant dense <2> : vector<4xi32>
235  %0 = math.fpowi %arg0, %c1 : f32, i32
236  %1 = math.fpowi %arg1, %v1 : vector<4xf32>, vector<4xi32>
237  %cm1 = arith.constant -2 : i32
238  %vm1 = arith.constant dense <-2> : vector<4xi32>
239  %2 = math.fpowi %arg0, %cm1 : f32, i32
240  %3 = math.fpowi %arg1, %vm1 : vector<4xf32>, vector<4xi32>
241  return %0, %1, %2, %3 : f32, vector<4xf32>, f32, vector<4xf32>
242}
243
244// CHECK-LABEL: @fpowi_exp_three(
245// CHECK-SAME: %[[ARG0:.+]]: f32
246// CHECK-SAME: %[[ARG1:.+]]: vector<4xf32>
247// CHECK-SAME: -> (f32, vector<4xf32>, f32, vector<4xf32>) {
248func.func @fpowi_exp_three(%arg0: f32, %arg1: vector<4xf32>) -> (f32, vector<4xf32>, f32, vector<4xf32>) {
249  // CHECK-DAG: %[[CST_S:.*]] = arith.constant 1.000000e+00 : f32
250  // CHECK-DAG: %[[CST_V:.*]] = arith.constant dense<1.000000e+00> : vector<4xf32>
251  // CHECK: %[[SMUL0:.*]] = arith.mulf %[[ARG0]], %[[ARG0]]
252  // CHECK: %[[SCALAR0:.*]] = arith.mulf %[[SMUL0]], %[[ARG0]]
253  // CHECK: %[[VMUL0:.*]] = arith.mulf %[[ARG1]], %[[ARG1]]
254  // CHECK: %[[VECTOR0:.*]] = arith.mulf %[[VMUL0]], %[[ARG1]]
255  // CHECK: %[[SCALAR1:.*]] = arith.divf %[[CST_S]], %[[ARG0]]
256  // CHECK: %[[SMUL1:.*]] = arith.mulf %[[SCALAR1]], %[[SCALAR1]]
257  // CHECK: %[[SMUL2:.*]] = arith.mulf %[[SMUL1]], %[[SCALAR1]]
258  // CHECK: %[[VECTOR1:.*]] = arith.divf %[[CST_V]], %[[ARG1]]
259  // CHECK: %[[VMUL1:.*]] = arith.mulf %[[VECTOR1]], %[[VECTOR1]]
260  // CHECK: %[[VMUL2:.*]] = arith.mulf %[[VMUL1]], %[[VECTOR1]]
261  // CHECK: return %[[SCALAR0]], %[[VECTOR0]], %[[SMUL2]], %[[VMUL2]]
262  %c1 = arith.constant 3 : i32
263  %v1 = arith.constant dense <3> : vector<4xi32>
264  %0 = math.fpowi %arg0, %c1 : f32, i32
265  %1 = math.fpowi %arg1, %v1 : vector<4xf32>, vector<4xi32>
266  %cm1 = arith.constant -3 : i32
267  %vm1 = arith.constant dense <-3> : vector<4xi32>
268  %2 = math.fpowi %arg0, %cm1 : f32, i32
269  %3 = math.fpowi %arg1, %vm1 : vector<4xf32>, vector<4xi32>
270  return %0, %1, %2, %3 : f32, vector<4xf32>, f32, vector<4xf32>
271}
272