[flang] Propagate fastmath flags during intrinsics simplification.

[vzakhari]

Summary

In general, the meaning of fastmath flags on a call during inlining
is that the call's operation flags must be ignored. For user functions
that means that the fastmath flags used for the function definition
override any call site's fastmath flags. For intrinsic functions
we can use the call site's fastmath flags, but we have to make sure
that the call sites with different flags produce/use different
simplified versions of the same intrinsic function.

[vzakhari]

Note that this only works for dot_product now. If it looks okay, I will extend it to other cases.

[jeanPerier]

Looks good to me. Have you considered having the flag copy part happening in fir::FirOpBuilder builder(op, kindMap); since the ctor has the op and is setting the builder ? But maybe that is too hidden and may not always be desired, so I am also OK with the explicit flag copy.

[vzakhari]

In D138048#3929942, @jeanPerier wrote:

Looks good to me. Have you considered having the flag copy part happening in fir::FirOpBuilder builder(op, kindMap); since the ctor has the op and is setting the builder ? But maybe that is too hidden and may not always be desired, so I am also OK with the explicit flag copy.

Thank you, Jean! Yes, I thought about automatically setting fastmath flags from the passed op, but it did not seem right to me. By definition, the passed op defines just the insertion point and does not necessarily refer to the operation that is being transformed. So there is a possibility that a pass transforms op1 with some fastmath flags, but uses op2 with other fastmath flags as an insertion point. I do not think it is possible now, but it may be the case later with MLIR-level inlining.

[vzakhari]

Added support for reduction cases.

[jeanPerier]

In D138048#3931051, @vzakhari wrote:

In D138048#3929942, @jeanPerier wrote:

Looks good to me. Have you considered having the flag copy part happening in fir::FirOpBuilder builder(op, kindMap); since the ctor has the op and is setting the builder ? But maybe that is too hidden and may not always be desired, so I am also OK with the explicit flag copy.

Thank you, Jean! Yes, I thought about automatically setting fastmath flags from the passed op, but it did not seem right to me. By definition, the passed op defines just the insertion point and does not necessarily refer to the operation that is being transformed. So there is a possibility that a pass transforms op1 with some fastmath flags, but uses op2 with other fastmath flags as an insertion point. I do not think it is possible now, but it may be the case later with MLIR-level inlining.

Ok, that makes sense to me.

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

public:
   void setFastMathFlags(mlir::arith::FastMathFlags flags) {
     fastMathFlags = flags;
   }
 
   /// Set default FastMathFlags value from the passed MathOptionsBase
   /// config.
   void setFastMathFlags(Fortran::common::MathOptionsBase options);
 
+  /// Get current FastMathFlags value.
+  mlir::arith::FastMathFlags getFastMathFlags() const { return fastMathFlags; }
+
   /// 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;
 

flang/lib/Optimizer/Transforms/SimplifyIntrinsics.cpp

private:
   /// \p kindMap is used to create FIROpBuilder
   /// \p genBodyFunc is the callback that builds the replacement function
   void simplifyReduction(fir::CallOp call, const fir::KindMapping &kindMap,
                          GenReductionBodyTy genBodyFunc);
 };
 
 } // namespace
 
+/// Create FirOpBuilder with the provided \p op insertion point
+/// and \p kindMap additionally inheriting FastMathFlags from \p op.
+static fir::FirOpBuilder
+getSimplificationBuilder(mlir::Operation *op, const fir::KindMapping &kindMap) {
+  fir::FirOpBuilder builder{op, kindMap};
+  auto fmi = mlir::dyn_cast<mlir::arith::ArithFastMathInterface>(*op);
+  if (!fmi)
+    return builder;
+
+  // Regardless of what default FastMathFlags are used by FirOpBuilder,
+  // override them with FastMathFlags attached to the operation.
+  builder.setFastMathFlags(fmi.getFastMathFlagsAttr().getValue());
+  return builder;
+}
+
+/// Stringify FastMathFlags set for the given \p builder in a way
+/// that the string may be used for mangling a function name.
+/// If FastMathFlags are set to 'none', then the result is an empty
+/// string.
+static std::string getFastMathFlagsString(const fir::FirOpBuilder &builder) {
+  mlir::arith::FastMathFlags flags = builder.getFastMathFlags();
+  if (flags == mlir::arith::FastMathFlags::none)
+    return {};
+
+  std::string fmfString{mlir::arith::stringifyFastMathFlags(flags)};
+  std::replace(fmfString.begin(), fmfString.end(), ',', '_');
+  return fmfString;
+}
+
 /// Generate function type for the simplified version of RTNAME(Sum) and
 /// similar functions with a fir.box<none> type returning \p elementType.
 static mlir::FunctionType genNoneBoxType(fir::FirOpBuilder &builder,
                                          const mlir::Type &elementType) {
   mlir::Type boxType = fir::BoxType::get(builder.getNoneType());
   return mlir::FunctionType::get(builder.getContext(), {boxType},
                                  {elementType});
 }
void SimplifyIntrinsicsPass::simplifyReduction(fir::CallOp call,
   const mlir::Value &dim = args[3];
   const mlir::Value &mask = args[4];
   // dim is zero when it is absent, which is an implementation
   // detail in the runtime library.
   bool dimAndMaskAbsent = isZero(dim) && isOperandAbsent(mask);
   unsigned rank = getDimCount(args[0]);
   if (dimAndMaskAbsent && rank > 0) {
     mlir::Location loc = call.getLoc();
-    fir::FirOpBuilder builder(call, kindMap);
+    fir::FirOpBuilder builder{getSimplificationBuilder(call, kindMap)};
+    std::string fmfString{getFastMathFlagsString(builder)};
 
     // Support only floating point and integer results now.
     mlir::Type resultType = call.getResult(0).getType();
     if (!resultType.isa<mlir::FloatType>() &&
         !resultType.isa<mlir::IntegerType>())
       return;
 
     auto argType = getArgElementType(args[0]);
     if (!argType)
       return;
     assert(*argType == resultType &&
            "Argument/result types mismatch in reduction");
 
     auto typeGenerator = [&resultType](fir::FirOpBuilder &builder) {
       return genNoneBoxType(builder, resultType);
     };
     auto bodyGenerator = [&rank, &genBodyFunc](fir::FirOpBuilder &builder,
                                                mlir::func::FuncOp &funcOp) {
       genBodyFunc(builder, funcOp, rank);
     };
     // Mangle the function name with the rank value as "x<rank>".
     std::string funcName =
         (mlir::Twine{callee.getLeafReference().getValue(), "x"} +
-         mlir::Twine{rank})
+         mlir::Twine{rank} +
+         // We must mangle the generated function name with FastMathFlags
+         // value.
+         (fmfString.empty() ? mlir::Twine{} : mlir::Twine{"_", fmfString}))
             .str();
     mlir::func::FuncOp newFunc =
         getOrCreateFunction(builder, funcName, typeGenerator, bodyGenerator);
     auto newCall =
         builder.create<fir::CallOp>(loc, newFunc, mlir::ValueRange{args[0]});
     call->replaceAllUsesWith(newCall.getResults());
     call->dropAllReferences();
     call->erase();
if (auto call = mlir::dyn_cast<fir::CallOp>(op)) {
         if (funcName.startswith(RTNAME_STRING(DotProduct))) {
           LLVM_DEBUG(llvm::dbgs() << "Handling " << funcName << "\n");
           LLVM_DEBUG(llvm::dbgs() << "Call operation:\n"; op->dump();
                      llvm::dbgs() << "\n");
           mlir::Operation::operand_range args = call.getArgs();
           const mlir::Value &v1 = args[0];
           const mlir::Value &v2 = args[1];
           mlir::Location loc = call.getLoc();
-          fir::FirOpBuilder builder(op, kindMap);
+          fir::FirOpBuilder builder{getSimplificationBuilder(op, kindMap)};
+          // Stringize the builder's FastMathFlags flags for mangling
+          // the generated function name.
+          std::string fmfString{getFastMathFlagsString(builder)};
 
           mlir::Type type = call.getResult(0).getType();
           if (!type.isa<mlir::FloatType>() && !type.isa<mlir::IntegerType>())
             return;
 
           // Try to find the element types of the boxed arguments.
           auto arg1Type = getArgElementType(v1);
           auto arg2Type = getArgElementType(v2);
if (auto call = mlir::dyn_cast<fir::CallOp>(op)) {
                                            mlir::func::FuncOp &funcOp) {
             genRuntimeDotBody(builder, funcOp, *arg1Type, *arg2Type);
           };
 
           // Suffix the function name with the element types
           // of the arguments.
           std::string typedFuncName(funcName);
           llvm::raw_string_ostream nameOS(typedFuncName);
-          nameOS << "_";
+          // We must mangle the generated function name with FastMathFlags
+          // value.
+          if (!fmfString.empty())
+            nameOS << '_' << fmfString;
+          nameOS << '_';
           arg1Type->print(nameOS);
-          nameOS << "_";
+          nameOS << '_';
           arg2Type->print(nameOS);
 
           mlir::func::FuncOp newFunc = getOrCreateFunction(
               builder, typedFuncName, typeGenerator, bodyGenerator);
           auto newCall = builder.create<fir::CallOp>(loc, newFunc,
                                                      mlir::ValueRange{v1, v2});
           call->replaceAllUsesWith(newCall.getResults());
           call->dropAllReferences();

flang/test/Transforms/simplifyintrinsics.fir

 fir.global linkonce @_QQcl.2E2F746573742E66393000 constant : !fir.char<1,11> {
   %0 = fir.string_lit "./test.f90\00"(11) : !fir.char<1,11>
   fir.has_value %0 : !fir.char<1,11>
 }
 
 // CHECK-NOT: call{{.*}}_FortranASumInteger8(
 // CHECK: call @_FortranASumInteger8x2_simplified(
 // CHECK-NOT: call{{.*}}_FortranASumInteger8(
+
+// -----
+
+func.func @dot_f32_contract_reassoc(%arg0: !fir.box<!fir.array<?xf32>> {fir.bindc_name = "a"}, %arg1: !fir.box<!fir.array<?xf32>> {fir.bindc_name = "b"}) -> f32 {
+  %0 = fir.alloca f32 {bindc_name = "dot", uniq_name = "_QFdotEdot"}
+  %1 = fir.address_of(@_QQcl.2E2F646F742E66393000) : !fir.ref<!fir.char<1,10>>
+  %c3_i32 = arith.constant 3 : i32
+  %2 = fir.convert %arg0 : (!fir.box<!fir.array<?xf32>>) -> !fir.box<none>
+  %3 = fir.convert %arg1 : (!fir.box<!fir.array<?xf32>>) -> !fir.box<none>
+  %4 = fir.convert %1 : (!fir.ref<!fir.char<1,10>>) -> !fir.ref<i8>
+  %5 = fir.call @_FortranADotProductReal4(%2, %3, %4, %c3_i32) fastmath<contract,reassoc> : (!fir.box<none>, !fir.box<none>, !fir.ref<i8>, i32) -> f32
+  fir.store %5 to %0 : !fir.ref<f32>
+  %6 = fir.load %0 : !fir.ref<f32>
+  return %6 : f32
+}
+
+func.func @dot_f32_fast(%arg0: !fir.box<!fir.array<?xf32>> {fir.bindc_name = "a"}, %arg1: !fir.box<!fir.array<?xf32>> {fir.bindc_name = "b"}) -> f32 {
+  %0 = fir.alloca f32 {bindc_name = "dot", uniq_name = "_QFdotEdot"}
+  %1 = fir.address_of(@_QQcl.2E2F646F742E66393000) : !fir.ref<!fir.char<1,10>>
+  %c3_i32 = arith.constant 3 : i32
+  %2 = fir.convert %arg0 : (!fir.box<!fir.array<?xf32>>) -> !fir.box<none>
+  %3 = fir.convert %arg1 : (!fir.box<!fir.array<?xf32>>) -> !fir.box<none>
+  %4 = fir.convert %1 : (!fir.ref<!fir.char<1,10>>) -> !fir.ref<i8>
+  %5 = fir.call @_FortranADotProductReal4(%2, %3, %4, %c3_i32) fastmath<fast> : (!fir.box<none>, !fir.box<none>, !fir.ref<i8>, i32) -> f32
+  fir.store %5 to %0 : !fir.ref<f32>
+  %6 = fir.load %0 : !fir.ref<f32>
+  return %6 : f32
+}
+
+func.func private @_FortranADotProductReal4(!fir.box<none>, !fir.box<none>, !fir.ref<i8>, i32) -> f32 attributes {fir.runtime}
+fir.global linkonce @_QQcl.2E2F646F742E66393000 constant : !fir.char<1,10> {
+  %0 = fir.string_lit "./dot.f90\00"(10) : !fir.char<1,10>
+  fir.has_value %0 : !fir.char<1,10>
+}
+
+// CHECK-LABEL: @dot_f32_contract_reassoc
+// CHECK: fir.call @_FortranADotProductReal4_reassoc_contract_f32_f32_simplified(%2, %3) fastmath<reassoc,contract>
+// CHECK-LABEL: @dot_f32_fast
+// CHECK: fir.call @_FortranADotProductReal4_fast_f32_f32_simplified(%2, %3) fastmath<fast>
+// CHECK-LABEL: func.func private @_FortranADotProductReal4_reassoc_contract_f32_f32_simplified
+// CHECK: arith.mulf %{{.*}}, %{{.*}} fastmath<reassoc,contract> : f32
+// CHECK: arith.addf %{{.*}}, %{{.*}} fastmath<reassoc,contract> : f32
+// CHECK-LABEL: func.func private @_FortranADotProductReal4_fast_f32_f32_simplified
+// CHECK: arith.mulf %{{.*}}, %{{.*}} fastmath<fast> : f32
+// CHECK: arith.addf %{{.*}}, %{{.*}} fastmath<fast> : f32
+
+// -----
+
+func.func @sum_1d_real_contract_reassoc(%arg0: !fir.ref<!fir.array<10xf64>> {fir.bindc_name = "a"}) -> f64 {
+  %c10 = arith.constant 10 : index
+  %0 = fir.alloca f64 {bindc_name = "sum_1d_real", uniq_name = "_QFsum_1d_realEsum_1d_real"}
+  %1 = fir.shape %c10 : (index) -> !fir.shape<1>
+  %2 = fir.embox %arg0(%1) : (!fir.ref<!fir.array<10xf64>>, !fir.shape<1>) -> !fir.box<!fir.array<10xf64>>
+  %3 = fir.absent !fir.box<i1>
+  %c0 = arith.constant 0 : index
+  %4 = fir.address_of(@_QQcl.2E2F6973756D5F352E66393000) : !fir.ref<!fir.char<1,13>>
+  %c5_i32 = arith.constant 5 : i32
+  %5 = fir.convert %2 : (!fir.box<!fir.array<10xf64>>) -> !fir.box<none>
+  %6 = fir.convert %4 : (!fir.ref<!fir.char<1,13>>) -> !fir.ref<i8>
+  %7 = fir.convert %c0 : (index) -> i32
+  %8 = fir.convert %3 : (!fir.box<i1>) -> !fir.box<none>
+  %9 = fir.call @_FortranASumReal8(%5, %6, %c5_i32, %7, %8) fastmath<contract,reassoc> : (!fir.box<none>, !fir.ref<i8>, i32, i32, !fir.box<none>) -> f64
+  fir.store %9 to %0 : !fir.ref<f64>
+  %10 = fir.load %0 : !fir.ref<f64>
+  return %10 : f64
+}
+
+func.func @sum_1d_real_fast(%arg0: !fir.ref<!fir.array<10xf64>> {fir.bindc_name = "a"}) -> f64 {
+  %c10 = arith.constant 10 : index
+  %0 = fir.alloca f64 {bindc_name = "sum_1d_real", uniq_name = "_QFsum_1d_realEsum_1d_real"}
+  %1 = fir.shape %c10 : (index) -> !fir.shape<1>
+  %2 = fir.embox %arg0(%1) : (!fir.ref<!fir.array<10xf64>>, !fir.shape<1>) -> !fir.box<!fir.array<10xf64>>
+  %3 = fir.absent !fir.box<i1>
+  %c0 = arith.constant 0 : index
+  %4 = fir.address_of(@_QQcl.2E2F6973756D5F352E66393000) : !fir.ref<!fir.char<1,13>>
+  %c5_i32 = arith.constant 5 : i32
+  %5 = fir.convert %2 : (!fir.box<!fir.array<10xf64>>) -> !fir.box<none>
+  %6 = fir.convert %4 : (!fir.ref<!fir.char<1,13>>) -> !fir.ref<i8>
+  %7 = fir.convert %c0 : (index) -> i32
+  %8 = fir.convert %3 : (!fir.box<i1>) -> !fir.box<none>
+  %9 = fir.call @_FortranASumReal8(%5, %6, %c5_i32, %7, %8) fastmath<fast> : (!fir.box<none>, !fir.ref<i8>, i32, i32, !fir.box<none>) -> f64
+  fir.store %9 to %0 : !fir.ref<f64>
+  %10 = fir.load %0 : !fir.ref<f64>
+  return %10 : f64
+}
+
+func.func private @_FortranASumReal8(!fir.box<none>, !fir.ref<i8>, i32, i32, !fir.box<none>) -> f64 attributes {fir.runtime}
+fir.global linkonce @_QQcl.2E2F6973756D5F352E66393000 constant : !fir.char<1,13> {
+  %0 = fir.string_lit "./isum_5.f90\00"(13) : !fir.char<1,13>
+  fir.has_value %0 : !fir.char<1,13>
+}
+
+// CHECK-LABEL: @sum_1d_real_contract_reassoc
+// CHECK: fir.call @_FortranASumReal8x1_reassoc_contract_simplified(%5) fastmath<reassoc,contract>
+// CHECK-LABEL: @sum_1d_real_fast
+// CHECK: fir.call @_FortranASumReal8x1_fast_simplified(%5) fastmath<fast>
+// CHECK-LABEL: func.func private @_FortranASumReal8x1_reassoc_contract_simplified
+// CHECK: arith.addf %{{.*}}, %{{.*}} fastmath<reassoc,contract> : f64
+// CHECK-LABEL: func.func private @_FortranASumReal8x1_fast_simplified
+// CHECK: arith.addf %{{.*}}, %{{.*}} fastmath<fast> : f64