[flang] Initial support for FastMathAttr setup in FirOpBuilder.

[vzakhari]

Summary

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.

[jeanPerier]

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.

Do you see any solutions for the clone issue (i.e., would it be possible to remove the listener while cloning by implementing a clone override that remove itself as a listener, call the parent clone method, and adds itself back ? ) ?
Or do you think there would be other solutions to deal with the fastMathAttr setting without using a listener, but that would take much more effort at that stage ?

[vzakhari]

In D137390#3908138, @jeanPerier wrote:

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.

Do you see any solutions for the clone issue (i.e., would it be possible to remove the listener while cloning by implementing a clone override that remove itself as a listener, call the parent clone method, and adds itself back ? ) ?
Or do you think there would be other solutions to deal with the fastMathAttr setting without using a listener, but that would take much more effort at that stage ?

i.e., would it be possible to remove the listener while cloning by implementing a clone override that remove itself as a listener, call the parent clone method, and adds itself back ?

This is a possibility, but unfortunately it will be unreliable as well. clone is not a virtual method of OpBuilder, so any redefinition that we add in FirOpBuilder will be only invoked if the method is invoked via an object with static FirOpBuilder type. If one does this via the base OpBuilder pointer our redefinition will be ignored.

I think the "default" attributes setup for new operations produced via OpBuilder should be a first class citizen of OpBuilder itself. Then we can define it as "set attributes defined by the current configuration of OpBuilder only for new instructions, i.e. methods that copy existing instructions (e.g. clone) just copy the instructions verbatim" - something along these lines. This requires some public discussion. I would like to proceed with this hacky approach so that we collect data about the issues along the road and sculpture the interfaces that we need to make this work for Flang. And in the meantime this allows us to have some support for fastmath, so that we can collect performance data showing the impact of proper fastmath handling. I am not sure if the clone issue in FirOpBuilder will affect Flang at all - I hope having this fastmath support will highlight cases where it may be an issue. When we have all things sorted out, we can specify what we need from OpBuilder.

[tschuett]

flang/include/flang/Optimizer/Builder/FIRBuilder.h

417

The virtual is not needed when you have override.

[kiranchandramohan]

LGTM. Please wait for @jeanPerier.

flang/include/flang/Optimizer/Builder/FIRBuilder.h

408

Would we need to change the fastmath setting in practice? As a design would it be good to not change the setting once a Builder is created?

flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp

532

Nit: A one-line explanation before each case might be helpful. For e.g, I have give below an explanation for each case. Feel free to skip if it will be verbose.
-> Test that the fast-math flag setting is none by default.
-> Check that setting the fast-math flag is honoured while creating an operation and is not impacted by the setting on a copy.
-> Check the setting of fast-math flags on an operation created with a copied builder.
-> When FIROpbuilder is copied the fast-math flags are retained

[vzakhari]

flang/include/flang/Optimizer/Builder/FIRBuilder.h

408

Yes, I think we may want to change it in some cases. I tried to come up with some examples in https://discourse.llvm.org/t/attaching-arith-fastmathattr-to-operations-in-flang/66366. Maybe they are not quite robust, but I would like to keep this possibility. FWIW, llvm::IRBuilderBase supports changing fastmath settings (and it is used) along with helper guards to revert to the previous state after local flags override.

flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp

532

Sure! Will do.

[vzakhari]

Added comments in the test.

[vzakhari]

Thank you for the review, Kiran and Jean!

flang/include/flang/Optimizer/Builder/FIRBuilder.h

 class BoxValue;
 
 //===----------------------------------------------------------------------===//
 // FirOpBuilder
 //===----------------------------------------------------------------------===//
 
 /// Extends the MLIR OpBuilder to provide methods for building common FIR
 /// patterns.
-class FirOpBuilder : public mlir::OpBuilder {
+class FirOpBuilder : public mlir::OpBuilder, public mlir::OpBuilder::Listener {
 public:
   explicit FirOpBuilder(mlir::Operation *op, const fir::KindMapping &kindMap)
-      : OpBuilder{op}, kindMap{kindMap} {}
+      : OpBuilder{op, /*listener=*/this}, kindMap{kindMap} {}
   explicit FirOpBuilder(mlir::OpBuilder &builder,
                         const fir::KindMapping &kindMap)
-      : OpBuilder{builder}, kindMap{kindMap} {}
+      : OpBuilder{builder}, kindMap{kindMap} {
+    setListener(this);
+  }
+
+  // The listener self-reference has to be updated in case of copy-construction.
+  FirOpBuilder(const FirOpBuilder &other)
+      : OpBuilder{other}, kindMap{other.kindMap}, fastMathFlags{
+                                                      other.fastMathFlags} {
+    setListener(this);
+  }
 
   /// Get the current Region of the insertion point.
   mlir::Region &getRegion() { return *getBlock()->getParent(); }
 
   /// Get the current Module
   mlir::ModuleOp getModule() {
     return getRegion().getParentOfType<mlir::ModuleOp>();
   }
public:
   mlir::Value genIsNullAddr(mlir::Location loc, mlir::Value addr);
 
   /// Compute the extent of (lb:ub:step) as max((ub-lb+step)/step, 0). See
   /// Fortran 2018 9.5.3.3.2 section for more details.
   mlir::Value genExtentFromTriplet(mlir::Location loc, mlir::Value lb,
                                    mlir::Value ub, mlir::Value step,
                                    mlir::Type type);
 
+  /// Set default FastMathFlags value for all operations
+  /// supporting mlir::arith::FastMathAttr that will be created
+  /// by this builder.
+  void setFastMathFlags(mlir::arith::FastMathFlags flags) {
+    fastMathFlags = flags;
+  }
+
   /// Dump the current function. (debug)
   LLVM_DUMP_METHOD void dumpFunc();
 
 private:
+  /// Set attributes (e.g. FastMathAttr) to \p op operation
+  /// based on the current attributes setting.
+  void setCommonAttributes(mlir::Operation *op) const;
+
+  /// FirOpBuilder hook for creating new operation.
+  void notifyOperationInserted(mlir::Operation *op) override {
+    setCommonAttributes(op);
+  }
+
   const KindMapping &kindMap;
+
+  /// FastMathFlags that need to be set for operations that support
+  /// mlir::arith::FastMathAttr.
+  mlir::arith::FastMathFlags fastMathFlags{};
 };
 
 } // namespace fir
 
 namespace fir::factory {
 
 //===----------------------------------------------------------------------===//
 // ExtendedValue inquiry helpers

flang/lib/Optimizer/Builder/FIRBuilder.cpp

mlir::Value fir::FirOpBuilder::genExtentFromTriplet(mlir::Location loc,
   auto diff = create<mlir::arith::SubIOp>(loc, ub, lb);
   auto add = create<mlir::arith::AddIOp>(loc, diff, step);
   auto div = create<mlir::arith::DivSIOp>(loc, add, step);
   auto cmp = create<mlir::arith::CmpIOp>(loc, mlir::arith::CmpIPredicate::sgt,
                                          div, zero);
   return create<mlir::arith::SelectOp>(loc, cmp, div, zero);
 }
 
+void fir::FirOpBuilder::setCommonAttributes(mlir::Operation *op) const {
+  auto fmi = mlir::dyn_cast<mlir::arith::ArithFastMathInterface>(*op);
+  if (!fmi)
+    return;
+  // TODO: use fmi.setFastMathFlagsAttr() after D137114 is merged.
+  //       For now set the attribute by the name.
+  llvm::StringRef arithFMFAttrName = fmi.getFastMathAttrName();
+  if (fastMathFlags != mlir::arith::FastMathFlags::none)
+    op->setAttr(arithFMFAttrName, mlir::arith::FastMathFlagsAttr::get(
+                                      op->getContext(), fastMathFlags));
+}
+
 //===--------------------------------------------------------------------===//
 // ExtendedValue inquiry helper implementation
 //===--------------------------------------------------------------------===//
 
 mlir::Value fir::factory::readCharLen(fir::FirOpBuilder &builder,
                                       mlir::Location loc,
                                       const fir::ExtendedValue &box) {
   return box.match(

flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp

for (const auto &scalar : derivedWithLengthScalars) {
     EXPECT_TRUE(fir::isDerivedWithLenParameters(scalar));
   }
   for (const auto &array : derivedWithLengthArrays) {
     EXPECT_EQ(fir::getBaseTypeOf(array), derivedWithLengthSeqTy);
     EXPECT_EQ(fir::getElementTypeOf(array), derivedWithLengthTy);
     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);
+}