[flang-commits] [flang] [flang][HLFIR] Add SeparateAllocatableAssign pass (PR #197814)

[Tom Eccles via flang-commits]

================
@@ -0,0 +1,173 @@
+//===- SeparateAllocatableAssign.cpp - Split realloc from assign ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+// Transform hlfir.assign with realloc semantics into a conditional
+// reallocation of the LHS followed by a plain hlfir.assign (without realloc).
+//
+// Before:
+//   hlfir.assign %rhs to %lhs realloc
+//
+// After:
+//   %shape = shape_of(%rhs)
+//   %new_lhs = genReallocIfNeeded(%lhs, %shape)  // host-side alloc
+//   hlfir.assign %rhs to %new_lhs                // element copy
+//
+// This is useful for OpenACC/OpenMP offloading where the allocation must
+// happen on the host before entering a device compute region.
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Builder/BoxValue.h"
+#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Builder/HLFIRTools.h"
+#include "flang/Optimizer/Builder/MutableBox.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/HLFIR/HLFIROps.h"
+#include "flang/Optimizer/HLFIR/Passes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/Support/Debug.h"
+
+namespace hlfir {
+#define GEN_PASS_DEF_SEPARATEALLOCATABLEASSIGN
+#include "flang/Optimizer/HLFIR/Passes.h.inc"
+} // namespace hlfir
+
+#define DEBUG_TYPE "separate-allocatable-assign"
+
+namespace {
+
+class SeparateAllocatableAssignConversion
+    : public mlir::OpRewritePattern<hlfir::AssignOp> {
+public:
+  using mlir::OpRewritePattern<hlfir::AssignOp>::OpRewritePattern;
+
+  llvm::LogicalResult
+  matchAndRewrite(hlfir::AssignOp assign,
+                  mlir::PatternRewriter &rewriter) const override {
+    if (!assign.isAllocatableAssignment())
+      return rewriter.notifyMatchFailure(assign, "not an allocatable assign");
+
+    hlfir::Entity rhs{assign.getRhs()};
+    hlfir::Entity lhs{assign.getLhs()};
+
+    if (!rhs.isArray())
+      return rewriter.notifyMatchFailure(assign, "RHS is not an array");
+
+    if (!lhs.isArray())
+      return rewriter.notifyMatchFailure(assign, "LHS is not an array");
+
+    mlir::Type rhsEleTy = rhs.getFortranElementType();
+    if (!fir::isa_trivial(rhsEleTy))
+      return rewriter.notifyMatchFailure(assign, "RHS type is not trivial");
+
+    mlir::Type lhsEleTy = lhs.getFortranElementType();
+    if (!fir::isa_trivial(lhsEleTy))
+      return rewriter.notifyMatchFailure(assign, "LHS type is not trivial");
+
+    if (lhsEleTy != rhsEleTy)
+      return rewriter.notifyMatchFailure(assign, "element type mismatch");
+
+    if (!fir::isBoxAddress(lhs.getType()))
+      return rewriter.notifyMatchFailure(assign, "LHS is not a box address");
+
+    // Reallocation frees the old LHS storage. If RHS references that same
+    // storage (e.g. a = a(:n)), the freed pointer would be read by the
+    // subsequent non-realloc assign, causing use-after-free.  Walk the
+    // RHS use-def chain; if any operand leads back to the LHS value,
+    // the two may alias and we must keep the original realloc assign.
+    if (!mlir::isa<hlfir::ExprType>(rhs.getType())) {
----------------
tblah wrote:

1) RHS may not always be a hlfir.expr type. If it isn't then there is still no check for lhs=rhs. Examples written and verified by codex:
```
  real, allocatable :: dst(:)
  real :: src(10)
  dst = src

  lowers to:

  hlfir.assign %src#0 to %dst#0 realloc
    : !fir.ref<!fir.array<10xf32>>, ...

  and a section assignment:

  dst = src(:n)

  lowers to a hlfir.designate result:

  hlfir.assign %section to %dst#0 realloc
    : !fir.box<!fir.array<?xf32>>, ...
```

2) Please use fir::AliasAnalysis to check if two memory locations may alias. There are plenty of non-trivial cases, pass through ("view-like") operations to handle, etc so it is better not to re-invent the wheel.

I will be offline for a couple of weeks after today. I trust the other reviewers on this patch to verify the fix so feel free to merge without another round of review from me if they approve.

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