Skip to content

Commit bcbf003

Browse files
meehatpameehatpa
committed
[mlir][spirv] Add atan and atan2 pattern to MathToSPIRV Conversion pass
Add missing math.atan to spirv.CL.atan and math.atan2 to spirv.CL.atan2 in MathToSPIRV.
1 parent 52220c2 commit bcbf003

File tree

2 files changed

+53
-39
lines changed

2 files changed

+53
-39
lines changed

‎mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -431,6 +431,8 @@ void populateMathToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
431431
patterns.add<Log1pOpPattern<spirv::CLLogOp>, ExpM1OpPattern<spirv::CLExpOp>,
432432
CheckedElementwiseOpPattern<math::AbsFOp, spirv::CLFAbsOp>,
433433
CheckedElementwiseOpPattern<math::AbsIOp, spirv::CLSAbsOp>,
434+
CheckedElementwiseOpPattern<math::AtanOp, spirv::CLAtanOp>,
435+
CheckedElementwiseOpPattern<math::Atan2Op, spirv::CLAtan2Op>,
434436
CheckedElementwiseOpPattern<math::CeilOp, spirv::CLCeilOp>,
435437
CheckedElementwiseOpPattern<math::CosOp, spirv::CLCosOp>,
436438
CheckedElementwiseOpPattern<math::ErfOp, spirv::CLErfOp>,

‎mlir/test/Conversion/MathToSPIRV/math-to-opencl-spirv.mlir

Lines changed: 51 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -4,83 +4,91 @@ module attributes { spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Kerne
44

55
// CHECK-LABEL: @float32_unary_scalar
66
func.func @float32_unary_scalar(%arg0: f32) {
7+
// CHECK: spirv.CL.atan %{{.*}}: f32
8+
%0 = math.atan %arg0 : f32
79
// CHECK: spirv.CL.cos %{{.*}}: f32
8-
%0 = math.cos %arg0 : f32
10+
%1 = math.cos %arg0 : f32
911
// CHECK: spirv.CL.exp %{{.*}}: f32
10-
%1 = math.exp %arg0 : f32
12+
%2 = math.exp %arg0 : f32
1113
// CHECK: %[[EXP:.+]] = spirv.CL.exp %arg0
1214
// CHECK: %[[ONE:.+]] = spirv.Constant 1.000000e+00 : f32
1315
// CHECK: spirv.FSub %[[EXP]], %[[ONE]]
14-
%2 = math.expm1 %arg0 : f32
16+
%3 = math.expm1 %arg0 : f32
1517
// CHECK: spirv.CL.log %{{.*}}: f32
16-
%3 = math.log %arg0 : f32
18+
%4 = math.log %arg0 : f32
1719
// CHECK: %[[ONE:.+]] = spirv.Constant 1.000000e+00 : f32
1820
// CHECK: %[[ADDONE:.+]] = spirv.FAdd %[[ONE]], %{{.+}}
1921
// CHECK: spirv.CL.log %[[ADDONE]]
20-
%4 = math.log1p %arg0 : f32
22+
%5 = math.log1p %arg0 : f32
2123
// CHECK: spirv.CL.rint %{{.*}}: f32
22-
%5 = math.roundeven %arg0 : f32
24+
%6 = math.roundeven %arg0 : f32
2325
// CHECK: spirv.CL.rsqrt %{{.*}}: f32
24-
%6 = math.rsqrt %arg0 : f32
26+
%7 = math.rsqrt %arg0 : f32
2527
// CHECK: spirv.CL.sqrt %{{.*}}: f32
26-
%7 = math.sqrt %arg0 : f32
28+
%8 = math.sqrt %arg0 : f32
2729
// CHECK: spirv.CL.tanh %{{.*}}: f32
28-
%8 = math.tanh %arg0 : f32
30+
%9 = math.tanh %arg0 : f32
2931
// CHECK: spirv.CL.sin %{{.*}}: f32
30-
%9 = math.sin %arg0 : f32
32+
%10 = math.sin %arg0 : f32
3133
// CHECK: spirv.CL.fabs %{{.*}}: f32
32-
%10 = math.absf %arg0 : f32
34+
%11 = math.absf %arg0 : f32
3335
// CHECK: spirv.CL.ceil %{{.*}}: f32
34-
%11 = math.ceil %arg0 : f32
36+
%12 = math.ceil %arg0 : f32
3537
// CHECK: spirv.CL.floor %{{.*}}: f32
36-
%12 = math.floor %arg0 : f32
38+
%13 = math.floor %arg0 : f32
3739
// CHECK: spirv.CL.erf %{{.*}}: f32
38-
%13 = math.erf %arg0 : f32
40+
%14 = math.erf %arg0 : f32
3941
// CHECK: spirv.CL.round %{{.*}}: f32
40-
%14 = math.round %arg0 : f32
42+
%15 = math.round %arg0 : f32
4143
return
4244
}
4345

4446
// CHECK-LABEL: @float32_unary_vector
4547
func.func @float32_unary_vector(%arg0: vector<3xf32>) {
48+
// CHECK: spirv.CL.atan %{{.*}}: vector<3xf32>
49+
%0 = math.atan %arg0 : vector<3xf32>
4650
// CHECK: spirv.CL.cos %{{.*}}: vector<3xf32>
47-
%0 = math.cos %arg0 : vector<3xf32>
51+
%1 = math.cos %arg0 : vector<3xf32>
4852
// CHECK: spirv.CL.exp %{{.*}}: vector<3xf32>
49-
%1 = math.exp %arg0 : vector<3xf32>
53+
%2 = math.exp %arg0 : vector<3xf32>
5054
// CHECK: %[[EXP:.+]] = spirv.CL.exp %arg0
5155
// CHECK: %[[ONE:.+]] = spirv.Constant dense<1.000000e+00> : vector<3xf32>
5256
// CHECK: spirv.FSub %[[EXP]], %[[ONE]]
53-
%2 = math.expm1 %arg0 : vector<3xf32>
57+
%3 = math.expm1 %arg0 : vector<3xf32>
5458
// CHECK: spirv.CL.log %{{.*}}: vector<3xf32>
55-
%3 = math.log %arg0 : vector<3xf32>
59+
%4 = math.log %arg0 : vector<3xf32>
5660
// CHECK: %[[ONE:.+]] = spirv.Constant dense<1.000000e+00> : vector<3xf32>
5761
// CHECK: %[[ADDONE:.+]] = spirv.FAdd %[[ONE]], %{{.+}}
5862
// CHECK: spirv.CL.log %[[ADDONE]]
59-
%4 = math.log1p %arg0 : vector<3xf32>
63+
%5 = math.log1p %arg0 : vector<3xf32>
6064
// CHECK: spirv.CL.rint %{{.*}}: vector<3xf32>
61-
%5 = math.roundeven %arg0 : vector<3xf32>
65+
%6 = math.roundeven %arg0 : vector<3xf32>
6266
// CHECK: spirv.CL.rsqrt %{{.*}}: vector<3xf32>
63-
%6 = math.rsqrt %arg0 : vector<3xf32>
67+
%7 = math.rsqrt %arg0 : vector<3xf32>
6468
// CHECK: spirv.CL.sqrt %{{.*}}: vector<3xf32>
65-
%7 = math.sqrt %arg0 : vector<3xf32>
69+
%8 = math.sqrt %arg0 : vector<3xf32>
6670
// CHECK: spirv.CL.tanh %{{.*}}: vector<3xf32>
67-
%8 = math.tanh %arg0 : vector<3xf32>
71+
%9 = math.tanh %arg0 : vector<3xf32>
6872
// CHECK: spirv.CL.sin %{{.*}}: vector<3xf32>
69-
%9 = math.sin %arg0 : vector<3xf32>
73+
%10 = math.sin %arg0 : vector<3xf32>
7074
return
7175
}
7276

7377
// CHECK-LABEL: @float32_binary_scalar
7478
func.func @float32_binary_scalar(%lhs: f32, %rhs: f32) {
79+
// CHECK: spirv.CL.atan2 %{{.*}}: f32
80+
%0 = math.atan2 %lhs, %rhs : f32
7581
// CHECK: spirv.CL.pow %{{.*}}: f32
76-
%0 = math.powf %lhs, %rhs : f32
82+
%1 = math.powf %lhs, %rhs : f32
7783
return
7884
}
7985

8086
// CHECK-LABEL: @float32_binary_vector
8187
func.func @float32_binary_vector(%lhs: vector<4xf32>, %rhs: vector<4xf32>) {
88+
// CHECK: spirv.CL.atan2 %{{.*}}: vector<4xf32>
89+
%0 = math.atan2 %lhs, %rhs : vector<4xf32>
8290
// CHECK: spirv.CL.pow %{{.*}}: vector<4xf32>
83-
%0 = math.powf %lhs, %rhs : vector<4xf32>
91+
%1 = math.powf %lhs, %rhs : vector<4xf32>
8492
return
8593
}
8694

@@ -118,18 +126,20 @@ module attributes {
118126

119127
// CHECK-LABEL: @vector_2d
120128
func.func @vector_2d(%arg0: vector<2x2xf32>) {
129+
// CHECK-NEXT: math.atan {{.+}} : vector<2x2xf32>
130+
%0 = math.atan %arg0 : vector<2x2xf32>
121131
// CHECK-NEXT: math.cos {{.+}} : vector<2x2xf32>
122-
%0 = math.cos %arg0 : vector<2x2xf32>
132+
%1 = math.cos %arg0 : vector<2x2xf32>
123133
// CHECK-NEXT: math.exp {{.+}} : vector<2x2xf32>
124-
%1 = math.exp %arg0 : vector<2x2xf32>
134+
%2 = math.exp %arg0 : vector<2x2xf32>
125135
// CHECK-NEXT: math.absf {{.+}} : vector<2x2xf32>
126-
%2 = math.absf %arg0 : vector<2x2xf32>
136+
%3 = math.absf %arg0 : vector<2x2xf32>
127137
// CHECK-NEXT: math.ceil {{.+}} : vector<2x2xf32>
128-
%3 = math.ceil %arg0 : vector<2x2xf32>
138+
%4 = math.ceil %arg0 : vector<2x2xf32>
129139
// CHECK-NEXT: math.floor {{.+}} : vector<2x2xf32>
130-
%4 = math.floor %arg0 : vector<2x2xf32>
140+
%5 = math.floor %arg0 : vector<2x2xf32>
131141
// CHECK-NEXT: math.powf {{.+}}, {{%.+}} : vector<2x2xf32>
132-
%5 = math.powf %arg0, %arg0 : vector<2x2xf32>
142+
%6 = math.powf %arg0, %arg0 : vector<2x2xf32>
133143
// CHECK-NEXT: return
134144
return
135145
}
@@ -138,18 +148,20 @@ func.func @vector_2d(%arg0: vector<2x2xf32>) {
138148

139149
// CHECK-LABEL: @tensor_1d
140150
func.func @tensor_1d(%arg0: tensor<2xf32>) {
151+
// CHECK-NEXT: math.atan {{.+}} : tensor<2xf32>
152+
%0 = math.atan %arg0 : tensor<2xf32>
141153
// CHECK-NEXT: math.cos {{.+}} : tensor<2xf32>
142-
%0 = math.cos %arg0 : tensor<2xf32>
154+
%1 = math.cos %arg0 : tensor<2xf32>
143155
// CHECK-NEXT: math.exp {{.+}} : tensor<2xf32>
144-
%1 = math.exp %arg0 : tensor<2xf32>
156+
%2 = math.exp %arg0 : tensor<2xf32>
145157
// CHECK-NEXT: math.absf {{.+}} : tensor<2xf32>
146-
%2 = math.absf %arg0 : tensor<2xf32>
158+
%3 = math.absf %arg0 : tensor<2xf32>
147159
// CHECK-NEXT: math.ceil {{.+}} : tensor<2xf32>
148-
%3 = math.ceil %arg0 : tensor<2xf32>
160+
%4 = math.ceil %arg0 : tensor<2xf32>
149161
// CHECK-NEXT: math.floor {{.+}} : tensor<2xf32>
150-
%4 = math.floor %arg0 : tensor<2xf32>
162+
%5 = math.floor %arg0 : tensor<2xf32>
151163
// CHECK-NEXT: math.powf {{.+}}, {{%.+}} : tensor<2xf32>
152-
%5 = math.powf %arg0, %arg0 : tensor<2xf32>
164+
%6 = math.powf %arg0, %arg0 : tensor<2xf32>
153165
// CHECK-NEXT: return
154166
return
155167
}

0 commit comments

Comments
 (0)