[flang-commits] [flang] [flang] Propagate fast-math flags to FIROpBuilder. (PR #126316)

[Slava Zakharin via flang-commits]

https://github.com/vzakhari updated https://github.com/llvm/llvm-project/pull/126316

>From 533163046fb1c4b0b075a7883ac0ff91450988e0 Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin at nvidia.com>
Date: Fri, 7 Feb 2025 13:56:00 -0800
Subject: [PATCH 1/2] [flang] Propagate fast-math flags to FIROpBuilder.

One constructor was missing to propagate fast-math flags
from an operation to the builder. It is fixed now.
---
 flang/include/flang/Optimizer/Builder/FIRBuilder.h | 8 +++++++-
 flang/test/HLFIR/maxval-elemental.fir              | 6 +++---
 flang/test/HLFIR/minval-elemental.fir              | 6 +++---
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/flang/include/flang/Optimizer/Builder/FIRBuilder.h b/flang/include/flang/Optimizer/Builder/FIRBuilder.h
index f2252d04079cf2c..64d27b0706f5867 100644
--- a/flang/include/flang/Optimizer/Builder/FIRBuilder.h
+++ b/flang/include/flang/Optimizer/Builder/FIRBuilder.h
@@ -57,7 +57,13 @@ class FirOpBuilder : public mlir::OpBuilder, public mlir::OpBuilder::Listener {
   explicit FirOpBuilder(mlir::Operation *op, fir::KindMapping kindMap,
                         mlir::SymbolTable *symbolTable = nullptr)
       : OpBuilder{op, /*listener=*/this}, kindMap{std::move(kindMap)},
-        symbolTable{symbolTable} {}
+        symbolTable{symbolTable} {
+    auto fmi = mlir::dyn_cast<mlir::arith::ArithFastMathInterface>(*op);
+    if (fmi) {
+      // Set the builder with FastMathFlags attached to the operation.
+      setFastMathFlags(fmi.getFastMathFlagsAttr().getValue());
+    }
+  }
   explicit FirOpBuilder(mlir::OpBuilder &builder, fir::KindMapping kindMap,
                         mlir::SymbolTable *symbolTable = nullptr)
       : OpBuilder(builder), OpBuilder::Listener(), kindMap{std::move(kindMap)},
diff --git a/flang/test/HLFIR/maxval-elemental.fir b/flang/test/HLFIR/maxval-elemental.fir
index aa642253b08323a..9dc028abe8da37a 100644
--- a/flang/test/HLFIR/maxval-elemental.fir
+++ b/flang/test/HLFIR/maxval-elemental.fir
@@ -82,9 +82,9 @@ func.func @_QPtest_float(%arg0: !fir.box<!fir.array<?xf32>> {fir.bindc_name = "a
 // CHECK-NEXT:  %[[V6:.*]] = fir.do_loop %arg1 = %c1 to %c4 step %c1 iter_args(%arg2 = %cst) -> (f32) {
 // CHECK-NEXT:      %[[V7:.*]] = hlfir.designate %[[V5]] (%arg1)  : (!fir.box<!fir.array<4xf32>>, index) -> !fir.ref<f32>
 // CHECK-NEXT:      %[[V8:.*]] = fir.load %[[V7]] : !fir.ref<f32>
-// CHECK-NEXT:      %[[V9:.*]] = arith.cmpf ogt, %[[V8]], %arg2 : f32
-// CHECK-NEXT:      %[[V10:.*]] = arith.cmpf une, %arg2, %arg2 : f32
-// CHECK-NEXT:      %[[V11:.*]] = arith.cmpf oeq, %[[V8]], %[[V8]] : f32
+// CHECK-NEXT:      %[[V9:.*]] = arith.cmpf ogt, %[[V8]], %arg2 fastmath<contract> : f32
+// CHECK-NEXT:      %[[V10:.*]] = arith.cmpf une, %arg2, %arg2 fastmath<contract> : f32
+// CHECK-NEXT:      %[[V11:.*]] = arith.cmpf oeq, %[[V8]], %[[V8]] fastmath<contract> : f32
 // CHECK-NEXT:      %[[V12:.*]] = arith.andi %[[V10]], %[[V11]] : i1
 // CHECK-NEXT:      %[[V13:.*]] = arith.ori %[[V9]], %[[V12]] : i1
 // CHECK-NEXT:      %[[V14:.*]] = arith.select %[[V13]], %[[V8]], %arg2 : f32
diff --git a/flang/test/HLFIR/minval-elemental.fir b/flang/test/HLFIR/minval-elemental.fir
index 8da1b1bdf515b9f..64cd5403ec558d4 100644
--- a/flang/test/HLFIR/minval-elemental.fir
+++ b/flang/test/HLFIR/minval-elemental.fir
@@ -82,9 +82,9 @@ func.func @_QPtest_float(%arg0: !fir.box<!fir.array<?xf32>> {fir.bindc_name = "a
 // CHECK-NEXT:  %[[V6:.*]] = fir.do_loop %arg1 = %c1 to %c4 step %c1 iter_args(%arg2 = %cst) -> (f32) {
 // CHECK-NEXT:      %[[V7:.*]] = hlfir.designate %[[V5]] (%arg1)  : (!fir.box<!fir.array<4xf32>>, index) -> !fir.ref<f32>
 // CHECK-NEXT:      %[[V8:.*]] = fir.load %[[V7]] : !fir.ref<f32>
-// CHECK-NEXT:      %[[V9:.*]] = arith.cmpf olt, %[[V8]], %arg2 : f32
-// CHECK-NEXT:      %[[V10:.*]] = arith.cmpf une, %arg2, %arg2 : f32
-// CHECK-NEXT:      %[[V11:.*]] = arith.cmpf oeq, %[[V8]], %[[V8]] : f32
+// CHECK-NEXT:      %[[V9:.*]] = arith.cmpf olt, %[[V8]], %arg2 fastmath<contract> : f32
+// CHECK-NEXT:      %[[V10:.*]] = arith.cmpf une, %arg2, %arg2 fastmath<contract> : f32
+// CHECK-NEXT:      %[[V11:.*]] = arith.cmpf oeq, %[[V8]], %[[V8]] fastmath<contract> : f32
 // CHECK-NEXT:      %[[V12:.*]] = arith.andi %[[V10]], %[[V11]] : i1
 // CHECK-NEXT:      %[[V13:.*]] = arith.ori %[[V9]], %[[V12]] : i1
 // CHECK-NEXT:      %[[V14:.*]] = arith.select %[[V13]], %[[V8]], %arg2 : f32

>From 667001577e4b16fb55992780780acab24e97abe7 Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin at nvidia.com>
Date: Fri, 7 Feb 2025 14:33:17 -0800
Subject: [PATCH 2/2] Changed FIROpBuilder creation in one opt-bufferization
 case.

---
 .../lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp b/flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp
index 0cfefc2d23ecb5d..3d506abbaa4542a 100644
--- a/flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp
+++ b/flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp
@@ -930,9 +930,7 @@ class ReductionConversion : public mlir::OpRewritePattern<Op> {
       llvm_unreachable("unsupported type");
     };
 
-    fir::KindMapping kindMap =
-        fir::getKindMapping(op->template getParentOfType<mlir::ModuleOp>());
-    fir::FirOpBuilder builder{op, kindMap};
+    fir::FirOpBuilder builder{rewriter, op.getOperation()};
 
     mlir::Value init;
     GenBodyFn genBodyFn;