diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp index 0543a624034ca..978120963c514 100644 --- a/flang/lib/Optimizer/Dialect/FIROps.cpp +++ b/flang/lib/Optimizer/Dialect/FIROps.cpp @@ -560,6 +560,16 @@ struct SimplifyArrayCoorOp : public mlir::OpRewritePattern { return mlir::failure(); } else if (auto reboxOp = mlir::dyn_cast_or_null( memref.getDefiningOp())) { + // Don't pull in rebox when the array_coor is inside an ACC construct + // and the rebox result 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 (op->getParentOfType() && + llvm::any_of(memref.getUsers(), [](mlir::Operation *u) { + return mlir::isa(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..5f6b8310a5ab9 --- /dev/null +++ b/flang/test/Fir/array-coor-canonicalization-acc.fir @@ -0,0 +1,48 @@ +// 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>, +// CHECK-SAME: %[[VAL:.*]]: f32) -> !fir.ref { +// 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]] +// CHECK: %[[COOR_OUT:.*]] = fir.array_coor %[[BOX]] %[[C1]] +// CHECK: return %[[COOR_OUT]] : !fir.ref +func.func @test_acc_copyin_preserves_rebox(%arg0: !fir.box>, %arg1: f32) -> !fir.ref { + %c1 = arith.constant 1 : index + %0 = fir.rebox %arg0 : (!fir.box>) -> !fir.box> + %1 = acc.copyin var(%0 : !fir.box>) -> !fir.box> {dataClause = #acc, name = "array"} + acc.parallel dataOperands(%1 : !fir.box>) { + %2 = fir.array_coor %0 %c1 : (!fir.box>, index) -> !fir.ref + fir.store %arg1 to %2 : !fir.ref + acc.yield + } attributes {defaultAttr = #acc} + acc.copyout accPtr(%1 : !fir.box>) to varPtr(%0 : !fir.box>) {dataClause = #acc, name = "array"} + %3 = fir.array_coor %0 %c1 : (!fir.box>, index) -> !fir.ref + return %3 : !fir.ref +} + +// ----- + +// 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.ref { +// CHECK: %[[C1:.*]] = arith.constant 1 : index +// CHECK: %[[COOR:.*]] = fir.array_coor %[[BOX]] %[[C1]] +// CHECK: return %[[COOR]] : !fir.ref +func.func @test_no_acc_folds_rebox(%arg0: !fir.box>) -> !fir.ref { + %c1 = arith.constant 1 : index + %0 = fir.rebox %arg0 : (!fir.box>) -> !fir.box> + %2 = fir.array_coor %0 %c1 : (!fir.box>, index) -> !fir.ref + return %2 : !fir.ref +}