[flang] Initial support for FastMathAttr setup in FirOpBuilder.

Provide FirOpBuilder::setFastMathFlags() to configure FastMathFlags
for the builder. Set FastMathAttr for operations based on FirOpBuilder
configuration via mlir::OpBuilder::Listener.

This is a little bit hacky solution, because we lose the ability
to hook other listeners to FirOpBuilder. There are also potential
issues with OpBuilder::clone() - the hook will be invoked for cloned
operations and will effectively overwrite FastMathAttr with the ones
configured in FirOpBuilder, which should not be happening.
We should teach mlir::OpBuilder about FastMathAttr setup in future.

Reviewed By: jeanPerier, kiranchandramohan

Differential Revision: https://reviews.llvm.org/D137390

1 files changed, 55 insertions, 0 deletions

diff --git a/flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp b/flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp
index 83d3defd3d06..9defe496b9c0 100644
--- a/flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp
+++ b/flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp
@@ -528,3 +528,58 @@ TEST_F(FIRBuilderTest, getBaseTypeOf) {
     EXPECT_TRUE(fir::isDerivedWithLenParameters(array));
   }
 }
+
+TEST_F(FIRBuilderTest, genArithFastMath) {
+  auto builder = getBuilder();
+  auto ctx = builder.getContext();
+  auto loc = builder.getUnknownLoc();
+
+  auto realTy = mlir::FloatType::getF32(ctx);
+  auto arg = builder.create<fir::UndefOp>(loc, realTy);
+
+  // Test that FastMathFlags is 'none' by default.
+  mlir::Operation *op1 = builder.create<mlir::arith::AddFOp>(loc, arg, arg);
+  auto op1_fmi =
+      mlir::dyn_cast_or_null<mlir::arith::ArithFastMathInterface>(op1);
+  EXPECT_TRUE(op1_fmi);
+  auto op1_fmf = op1_fmi.getFastMathFlagsAttr().getValue();
+  EXPECT_EQ(op1_fmf, arith::FastMathFlags::none);
+
+  // Test that the builder is copied properly.
+  fir::FirOpBuilder builder_copy(builder);
+
+  arith::FastMathFlags FMF1 =
+      arith::FastMathFlags::contract | arith::FastMathFlags::reassoc;
+  builder.setFastMathFlags(FMF1);
+  arith::FastMathFlags FMF2 =
+      arith::FastMathFlags::nnan | arith::FastMathFlags::ninf;
+  builder_copy.setFastMathFlags(FMF2);
+
+  // Modifying FastMathFlags for the copy must not affect the original builder.
+  mlir::Operation *op2 = builder.create<mlir::arith::AddFOp>(loc, arg, arg);
+  auto op2_fmi =
+      mlir::dyn_cast_or_null<mlir::arith::ArithFastMathInterface>(op2);
+  EXPECT_TRUE(op2_fmi);
+  auto op2_fmf = op2_fmi.getFastMathFlagsAttr().getValue();
+  EXPECT_EQ(op2_fmf, FMF1);
+
+  // Modifying FastMathFlags for the original builder must not affect the copy.
+  mlir::Operation *op3 =
+      builder_copy.create<mlir::arith::AddFOp>(loc, arg, arg);
+  auto op3_fmi =
+      mlir::dyn_cast_or_null<mlir::arith::ArithFastMathInterface>(op3);
+  EXPECT_TRUE(op3_fmi);
+  auto op3_fmf = op3_fmi.getFastMathFlagsAttr().getValue();
+  EXPECT_EQ(op3_fmf, FMF2);
+
+  // Test that the builder copy inherits FastMathFlags from the original.
+  fir::FirOpBuilder builder_copy2(builder);
+
+  mlir::Operation *op4 =
+      builder_copy2.create<mlir::arith::AddFOp>(loc, arg, arg);
+  auto op4_fmi =
+      mlir::dyn_cast_or_null<mlir::arith::ArithFastMathInterface>(op4);
+  EXPECT_TRUE(op4_fmi);
+  auto op4_fmf = op4_fmi.getFastMathFlagsAttr().getValue();
+  EXPECT_EQ(op4_fmf, FMF1);
+}