[flang-commits] [flang] [flang] Try to keep overflow flags in `FIRToMemRef::canonicalizeIndex'. (PR #220127)

[via flang-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-fir-hlfir

Author: Valery Dmitriev (valerydmit)

<details>
<summary>Changes</summary>


The method specifically runs on every `fir.array_coor` subscript. For `arith.addi` it recursively canonicalizes both operands and then unconditionally rebuilds the op defaulting flags to `none`. Thus subscript operand becomes a flagless clone, with the original one left behind dead. The patch makes a couple of changes: (1) use the original op if unmodified and (2) when its type preserved keep overflow flags of the original.

Assisted-by: Claude Opus 5

---
Full diff: https://github.com/llvm/llvm-project/pull/220127.diff


3 Files Affected:

- (modified) flang/lib/Optimizer/Transforms/FIRToMemRef.cpp (+16-5) 
- (added) flang/test/Transforms/FIRToMemRef/index-overflow-flags.mlir (+54) 
- (modified) flang/test/Transforms/FIRToMemRef/slice.mlir (+1-2) 


``````````diff
diff --git a/flang/lib/Optimizer/Transforms/FIRToMemRef.cpp b/flang/lib/Optimizer/Transforms/FIRToMemRef.cpp
index 5851022b1bc47..3975cddeed75b 100644
--- a/flang/lib/Optimizer/Transforms/FIRToMemRef.cpp
+++ b/flang/lib/Optimizer/Transforms/FIRToMemRef.cpp
@@ -1333,10 +1333,12 @@ FIRToMemRef::getFIRConvert(Operation *memOp, Operation *op,
 ///      `index_cast` if they need an index-typed result.
 ///
 ///   4. `arith.addi %a, %b`
-///      -> recursively canonicalize both operands, and if their result
-///      types match, build a new `arith.addi` at the same location. If
-///      the canonicalized operand types diverge, returns the original op
-///      untouched (the caller can still `index_cast` externally).
+///      -> recursively canonicalize both operands. If neither operand
+///      changed, the original op is returned as-is. Otherwise, if the
+///      canonicalized operand types match, build a new `arith.addi` at the
+///      same location, carrying over the original overflow flags.
+///      If the canonicalized operand types diverge, returns the original
+///      op untouched (the caller can still `index_cast` externally).
 ///
 /// Only these four patterns fire -- this is intentionally a narrow peephole,
 /// not a general folder. Multiplication, sub, cast chains through other ops,
@@ -1372,8 +1374,17 @@ Value FIRToMemRef::canonicalizeIndex(Value index,
   if (auto add = dyn_cast<arith::AddIOp>(op)) {
     Value lhs = canonicalizeIndex(add.getLhs(), rewriter);
     Value rhs = canonicalizeIndex(add.getRhs(), rewriter);
-    if (lhs.getType() == rhs.getType())
+    // Neither operand simplified, so a rebuilt op would be an exact duplicate.
+    // Reuse the original instead.
+    if (lhs == add.getLhs() && rhs == add.getRhs())
+      return index;
+    if (lhs.getType() == rhs.getType()) {
+      // Carry over overflow flags when width unchanged only.
+      if (lhs.getType() == add.getType())
+        return arith::AddIOp::create(rewriter, op->getLoc(), lhs, rhs,
+                                     add.getOverflowFlags());
       return arith::AddIOp::create(rewriter, op->getLoc(), lhs, rhs);
+    }
   }
   return index;
 }
diff --git a/flang/test/Transforms/FIRToMemRef/index-overflow-flags.mlir b/flang/test/Transforms/FIRToMemRef/index-overflow-flags.mlir
new file mode 100644
index 0000000000000..7e1befd8121a8
--- /dev/null
+++ b/flang/test/Transforms/FIRToMemRef/index-overflow-flags.mlir
@@ -0,0 +1,54 @@
+// Test that overflow flags on subscript arithmetic survive the index
+// canonicalization performed while lowering fir.array_coor.
+
+// RUN: fir-opt %s --fir-to-memref --allow-unregistered-dialect | FileCheck %s
+
+// An `arith.addi` whose operands need no canonicalization must be reused
+// as-is, keeping `overflow<nsw>`.
+//
+// CHECK-LABEL: func.func @addi_nsw_preserved
+// CHECK:       [[I:%.+]]    = memref.load
+// CHECK:       [[J:%.+]]    = memref.load
+// CHECK:       [[SUM:%.+]]  = arith.addi [[I]], [[J]] overflow<nsw> : i32
+// CHECK-NOT:   arith.addi {{.*}} : i32
+// CHECK:       [[CAST:%.+]] = arith.index_cast [[SUM]] : i32 to index
+func.func @addi_nsw_preserved(%arg0: !fir.ref<!fir.array<100xf32>>, %arg1: !fir.ref<i32>, %arg2: !fir.ref<i32>) {
+  %c100 = arith.constant 100 : index
+  %dscope = fir.undefined !fir.dscope
+  %shape = fir.shape %c100 : (index) -> !fir.shape<1>
+  %a = fir.declare %arg0(%shape) dummy_scope %dscope {uniq_name = "a"} : (!fir.ref<!fir.array<100xf32>>, !fir.shape<1>, !fir.dscope) -> !fir.ref<!fir.array<100xf32>>
+  %i = fir.declare %arg1 dummy_scope %dscope {uniq_name = "i"} : (!fir.ref<i32>, !fir.dscope) -> !fir.ref<i32>
+  %j = fir.declare %arg2 dummy_scope %dscope {uniq_name = "j"} : (!fir.ref<i32>, !fir.dscope) -> !fir.ref<i32>
+  %iv = fir.load %i : !fir.ref<i32>
+  %jv = fir.load %j : !fir.ref<i32>
+  %sum = arith.addi %iv, %jv overflow<nsw> : i32
+  %addr = fir.array_coor %a(%shape) %sum : (!fir.ref<!fir.array<100xf32>>, !fir.shape<1>, i32) -> !fir.ref<f32>
+  %v = fir.load %addr : !fir.ref<f32>
+  return
+}
+
+// When canonicalization peels the `arith.extsi` off both operands the add is
+// rebuilt at the narrower width, so `overflow<nsw>` must not carry over: a
+// 64-bit add that cannot wrap can still wrap in 32 bits.
+//
+// CHECK-LABEL: func.func @addi_nsw_dropped_on_narrowing
+// CHECK:       [[I:%.+]]    = memref.load
+// CHECK:       [[J:%.+]]    = memref.load
+// CHECK:       [[SUM:%.+]]  = arith.addi [[I]], [[J]] : i32
+// CHECK:       [[CAST:%.+]] = arith.index_cast [[SUM]] : i32 to index
+func.func @addi_nsw_dropped_on_narrowing(%arg0: !fir.ref<!fir.array<100xf32>>, %arg1: !fir.ref<i32>, %arg2: !fir.ref<i32>) {
+  %c100 = arith.constant 100 : index
+  %dscope = fir.undefined !fir.dscope
+  %shape = fir.shape %c100 : (index) -> !fir.shape<1>
+  %a = fir.declare %arg0(%shape) dummy_scope %dscope {uniq_name = "a"} : (!fir.ref<!fir.array<100xf32>>, !fir.shape<1>, !fir.dscope) -> !fir.ref<!fir.array<100xf32>>
+  %i = fir.declare %arg1 dummy_scope %dscope {uniq_name = "i"} : (!fir.ref<i32>, !fir.dscope) -> !fir.ref<i32>
+  %j = fir.declare %arg2 dummy_scope %dscope {uniq_name = "j"} : (!fir.ref<i32>, !fir.dscope) -> !fir.ref<i32>
+  %iv = fir.load %i : !fir.ref<i32>
+  %jv = fir.load %j : !fir.ref<i32>
+  %ie = arith.extsi %iv : i32 to i64
+  %je = arith.extsi %jv : i32 to i64
+  %sum = arith.addi %ie, %je overflow<nsw> : i64
+  %addr = fir.array_coor %a(%shape) %sum : (!fir.ref<!fir.array<100xf32>>, !fir.shape<1>, i64) -> !fir.ref<f32>
+  %v = fir.load %addr : !fir.ref<f32>
+  return
+}
diff --git a/flang/test/Transforms/FIRToMemRef/slice.mlir b/flang/test/Transforms/FIRToMemRef/slice.mlir
index 737babd4733a0..a1e0fba75aaf4 100644
--- a/flang/test/Transforms/FIRToMemRef/slice.mlir
+++ b/flang/test/Transforms/FIRToMemRef/slice.mlir
@@ -297,8 +297,7 @@ func.func @extract_column(%arg0: !fir.ref<!fir.array<100x5xf32>> {fir.bindc_name
 // CHECK:         [[BOXADDR2:%.*]] = fir.box_addr [[EMBOX]] : (!fir.box<!fir.array<7xf32>>) -> !fir.ref<!fir.array<7xf32>>
 // CHECK:         [[CONVERT:%.*]] = fir.convert [[BOXADDR2]] : (!fir.ref<!fir.array<7xf32>>) -> memref<7xf32>
 // CHECK:         %[[C1_0:.*]] = arith.constant 1 : index
-// CHECK:         [[ADD2:%.*]] = arith.addi %[[ARG0]], %[[C_NEG1]] : index
-// CHECK:         [[SUB1:%.*]] = arith.subi [[ADD2]], %[[C0]] : index
+// CHECK:         [[SUB1:%.*]] = arith.subi [[ADD1]], %[[C0]] : index
 // CHECK:         [[MUL1:%.*]] = arith.muli [[SUB1]], %[[C1_0]] : index
 // CHECK:         [[SUB2:%[0-9]+]] = arith.subi %[[C0]], %[[C0]] : index
 // CHECK:         [[ADD3:%.*]] = arith.addi [[MUL1]], [[SUB2]] : index

``````````

</details>


https://github.com/llvm/llvm-project/pull/220127