[flang-commits] [flang] [flang][openacc] Prevent SimplifyArrayCoorOp from folding rebox used by ACC data entry ops (PR #187616)

[Valentin Clement バレンタイン クレメン via flang-commits]

https://github.com/clementval created https://github.com/llvm/llvm-project/pull/187616

The SimplifyArrayCoorOp canonicalization pattern folds fir.rebox into
fir.array_coor by replacing the rebox result with its source. When the
rebox result is the variable of an acc.copyin (or other data entry op),
this folding breaks the chain that LegalizeDataValuesInRegion relies on
to replace host values with device pointers inside compute regions.
With default(none), ACCImplicitData cannot recover the mapping, leaving
the rebox source as an illegal live-in in the offload region.

Guard the folding so it bails out when any user of the rebox is an
ACC_DATA_ENTRY_OPS.

>From 1749e35c463bb7585e34280d991d580615aa3dc9 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Thu, 19 Mar 2026 17:51:47 -0700
Subject: [PATCH] [flang][openacc] Prevent SimplifyArrayCoorOp from folding
 rebox used by ACC data entry ops

---
 flang/lib/Optimizer/Dialect/FIROps.cpp        |  8 ++++
 .../Fir/array-coor-canonicalization-acc.fir   | 45 +++++++++++++++++++
 2 files changed, 53 insertions(+)
 create mode 100644 flang/test/Fir/array-coor-canonicalization-acc.fir

diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index 0543a624034ca..0f14355fcdeae 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -560,6 +560,14 @@ struct SimplifyArrayCoorOp : public mlir::OpRewritePattern<fir::ArrayCoorOp> {
         return mlir::failure();
     } else if (auto reboxOp = mlir::dyn_cast_or_null<fir::ReboxOp>(
                    memref.getDefiningOp())) {
+      // Don't pull in rebox when it is referenced by an ACC data clause.
+      // The data legalization pipeline relies on the rebox result being the
+      // copyin var; folding through it would leave the rebox source as an
+      // unhandled live-in inside the compute region.
+      if (llvm::any_of(memref.getUsers(), [](mlir::Operation *u) {
+            return mlir::isa<ACC_DATA_ENTRY_OPS>(u);
+          }))
+        return mlir::failure();
       boxedMemref = reboxOp.getBox();
       boxedShape = reboxOp.getShape();
       // Avoid pulling in rebox that performs reshaping.
diff --git a/flang/test/Fir/array-coor-canonicalization-acc.fir b/flang/test/Fir/array-coor-canonicalization-acc.fir
new file mode 100644
index 0000000000000..46017355b32fa
--- /dev/null
+++ b/flang/test/Fir/array-coor-canonicalization-acc.fir
@@ -0,0 +1,45 @@
+// RUN: fir-opt --canonicalize --split-input-file %s | FileCheck %s
+
+// Verify that fir.rebox is NOT folded into fir.array_coor when the rebox
+// result is used by an acc.copyin. The data legalization pipeline depends on
+// the copyin's var (the rebox result) matching the value used inside compute
+// regions; folding would break this chain and produce an illegal live-in.
+// CHECK-LABEL:   func.func @test_acc_copyin_preserves_rebox(
+// CHECK-SAME:        %[[BOX:.*]]: !fir.box<!fir.array<?xf32>>,
+// CHECK-SAME:        %[[VAL:.*]]: f32) {
+// CHECK:           %[[C1:.*]] = arith.constant 1 : index
+// CHECK:           %[[REBOX:.*]] = fir.rebox %[[BOX]]
+// CHECK:           %[[COPYIN:.*]] = acc.copyin var(%[[REBOX]]
+// CHECK:           acc.parallel dataOperands(%[[COPYIN]] :
+// CHECK:             %[[COOR:.*]] = fir.array_coor %[[REBOX]] %[[C1]]
+// CHECK:             fir.store %[[VAL]] to %[[COOR]]
+// CHECK:             acc.yield
+// CHECK:           }
+// CHECK:           acc.copyout accVar(%[[COPYIN]]
+func.func @test_acc_copyin_preserves_rebox(%arg0: !fir.box<!fir.array<?xf32>>, %arg1: f32) {
+  %c1 = arith.constant 1 : index
+  %0 = fir.rebox %arg0 : (!fir.box<!fir.array<?xf32>>) -> !fir.box<!fir.array<?xf32>>
+  %1 = acc.copyin var(%0 : !fir.box<!fir.array<?xf32>>) -> !fir.box<!fir.array<?xf32>> {dataClause = #acc<data_clause acc_copy>, name = "array"}
+  acc.parallel dataOperands(%1 : !fir.box<!fir.array<?xf32>>) {
+    %2 = fir.array_coor %0 %c1 : (!fir.box<!fir.array<?xf32>>, index) -> !fir.ref<f32>
+    fir.store %arg1 to %2 : !fir.ref<f32>
+    acc.yield
+  } attributes {defaultAttr = #acc<defaultvalue none>}
+  acc.copyout accPtr(%1 : !fir.box<!fir.array<?xf32>>) to varPtr(%0 : !fir.box<!fir.array<?xf32>>) {dataClause = #acc<data_clause acc_copy>, name = "array"}
+  return
+}
+
+// -----
+
+// Same scenario but without acc.copyin: rebox SHOULD be folded normally.
+// CHECK-LABEL:   func.func @test_no_acc_folds_rebox(
+// CHECK-SAME:        %[[BOX:.*]]: !fir.box<!fir.array<?xf32>>) -> !fir.ref<f32> {
+// CHECK:           %[[C1:.*]] = arith.constant 1 : index
+// CHECK:           %[[COOR:.*]] = fir.array_coor %[[BOX]] %[[C1]]
+// CHECK:           return %[[COOR]] : !fir.ref<f32>
+func.func @test_no_acc_folds_rebox(%arg0: !fir.box<!fir.array<?xf32>>) -> !fir.ref<f32> {
+  %c1 = arith.constant 1 : index
+  %0 = fir.rebox %arg0 : (!fir.box<!fir.array<?xf32>>) -> !fir.box<!fir.array<?xf32>>
+  %2 = fir.array_coor %0 %c1 : (!fir.box<!fir.array<?xf32>>, index) -> !fir.ref<f32>
+  return %2 : !fir.ref<f32>
+}