[mlir][complex] Convert complex.abs to arith with fastmath flag

[Lewuathe]

Summary

This is the progressive follow-up to support the conversion of the complex.abs to arith with fastmath flag .

See: https://reviews.llvm.org/rG14de5a2a4f27ed719b1f59f99d2730624a26f706

[kiranchandramohan]

Thanks. LGTM.

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

726

Nit: Newline

[Lewuathe]

It has been merged. (Sorry, I failed to add the link to this diff)
https://github.com/llvm/llvm-project/commit/653f77690bb2473ff405ccc681c73217fae8b748

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

struct AbsOpConversion : public OpConversionPattern<complex::AbsOp> {
   using OpConversionPattern<complex::AbsOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(complex::AbsOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
     auto type = op.getType();
 
+    arith::FastMathFlagsAttr fmf = op.getFastMathFlagsAttr();
+
     Value real =
         rewriter.create<complex::ReOp>(loc, type, adaptor.getComplex());
     Value imag =
         rewriter.create<complex::ImOp>(loc, type, adaptor.getComplex());
-    Value realSqr = rewriter.create<arith::MulFOp>(loc, real, real);
-    Value imagSqr = rewriter.create<arith::MulFOp>(loc, imag, imag);
-    Value sqNorm = rewriter.create<arith::AddFOp>(loc, realSqr, imagSqr);
+    Value realSqr =
+        rewriter.create<arith::MulFOp>(loc, real, real, fmf.getValue());
+    Value imagSqr =
+        rewriter.create<arith::MulFOp>(loc, imag, imag, fmf.getValue());
+    Value sqNorm =
+        rewriter.create<arith::AddFOp>(loc, realSqr, imagSqr, fmf.getValue());
 
     rewriter.replaceOpWithNewOp<math::SqrtOp>(op, sqNorm);
     return success();
   }
 };
 
 // atan2(y,x) = -i * log((x + i * y)/sqrt(x**2+y**2))
 struct Atan2OpConversion : public OpConversionPattern<complex::Atan2Op> {

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

 func.func @complex_angle(%arg: complex<f32>) -> f32 {
   %angle = complex.angle %arg : complex<f32>
   return %angle : f32
 }
 // CHECK: %[[REAL:.*]] = complex.re %[[ARG]] : complex<f32>
 // CHECK: %[[IMAG:.*]] = complex.im %[[ARG]] : complex<f32>
 // CHECK: %[[RESULT:.*]] = math.atan2 %[[IMAG]], %[[REAL]] : f32
 // CHECK: return %[[RESULT]] : f32
+
+// -----
+
+// CHECK-LABEL: func @complex_abs_with_fmf
+// CHECK-SAME: %[[ARG:.*]]: complex<f32>
+func.func @complex_abs_with_fmf(%arg: complex<f32>) -> f32 {
+  %abs = complex.abs %arg fastmath<nnan,contract> : complex<f32>
+  return %abs : f32
+}
+// CHECK: %[[REAL:.*]] = complex.re %[[ARG]] : complex<f32>
+// CHECK: %[[IMAG:.*]] = complex.im %[[ARG]] : complex<f32>
+// CHECK-DAG: %[[REAL_SQ:.*]] = arith.mulf %[[REAL]], %[[REAL]] fastmath<nnan,contract> : f32
+// CHECK-DAG: %[[IMAG_SQ:.*]] = arith.mulf %[[IMAG]], %[[IMAG]] fastmath<nnan,contract> : f32
+// CHECK: %[[SQ_NORM:.*]] = arith.addf %[[REAL_SQ]], %[[IMAG_SQ]] fastmath<nnan,contract> : f32
+// CHECK: %[[NORM:.*]] = math.sqrt %[[SQ_NORM]] : f32
+// CHECK: return %[[NORM]] : f32
 No newline at end of file