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

[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())) {
----------------
jeanPerier wrote:

I disagree with codex here in the sense that as long as the deallocation is not moved above the hlfir.expr producer, it is a problem of the hlfir.assign lowering/optimization to detect that it can move the hlfir.expr evaluation into the assignment.

For instance this should be done [by this code](https://github.com/llvm/llvm-project/blob/8466750c230d5e51d509a8071c8a50e7de390029/flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp#L388) for instance for hlfir.elemental inlining into the assignment (arguably this analysis may currently ignoring the allocation/deallocation effects which is bad).

You probably want to double check that this code is correctly not ignoring mem free effects, and if so I think you can just focus on analyzing variable RHS and leave it up to the expression bufferization to not move the hlfir.expr evaluation over deallocation.

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