Skip to content

Commit b5d3e0b

Browse files
committed
[mlir][complex] Fastmath flag support for complex.angle
1 parent c9ae553 commit b5d3e0b

File tree

2 files changed

+16
-2
lines changed

2 files changed

+16
-2
lines changed

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1143,13 +1143,14 @@ struct AngleOpConversion : public OpConversionPattern<complex::AngleOp> {
11431143
ConversionPatternRewriter &rewriter) const override {
11441144
auto loc = op.getLoc();
11451145
auto type = op.getType();
1146+
arith::FastMathFlagsAttr fmf = op.getFastMathFlagsAttr();
11461147

11471148
Value real =
11481149
rewriter.create<complex::ReOp>(loc, type, adaptor.getComplex());
11491150
Value imag =
11501151
rewriter.create<complex::ImOp>(loc, type, adaptor.getComplex());
11511152

1152-
rewriter.replaceOpWithNewOp<math::Atan2Op>(op, imag, real);
1153+
rewriter.replaceOpWithNewOp<math::Atan2Op>(op, imag, real, fmf);
11531154

11541155
return success();
11551156
}

mlir/test/Conversion/ComplexToStandard/convert-to-standard.mlir

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2035,4 +2035,17 @@ func.func @complex_tanh_with_fmf(%arg: complex<f32>) -> complex<f32> {
20352035
// CHECK: %[[NUM:.*]] = complex.create %[[TANH_A]], %[[TAN_B]] : complex<f32>
20362036
// CHECK: %[[ONE:.*]] = arith.constant 1.000000e+00 : f32
20372037
// CHECK: %[[MUL:.*]] = arith.mulf %[[TANH_A]], %[[TAN_B]] fastmath<nnan,contract> : f32
2038-
// CHECK: %[[DENOM:.*]] = complex.create %[[ONE]], %[[MUL]] : complex<f32>
2038+
// CHECK: %[[DENOM:.*]] = complex.create %[[ONE]], %[[MUL]] : complex<f32>
2039+
2040+
// -----
2041+
2042+
// CHECK-LABEL: func.func @complex_angle_with_fmf
2043+
// CHECK-SAME: %[[ARG:.*]]: complex<f32>
2044+
func.func @complex_angle_with_fmf(%arg: complex<f32>) -> f32 {
2045+
%angle = complex.angle %arg fastmath<nnan,contract> : complex<f32>
2046+
return %angle : f32
2047+
}
2048+
// CHECK: %[[REAL:.*]] = complex.re %[[ARG]] : complex<f32>
2049+
// CHECK: %[[IMAG:.*]] = complex.im %[[ARG]] : complex<f32>
2050+
// CHECK: %[[RESULT:.*]] = math.atan2 %[[IMAG]], %[[REAL]] fastmath<nnan,contract> : f32
2051+
// CHECK: return %[[RESULT]] : f32

0 commit comments

Comments
 (0)