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

[via flang-commits]

================
@@ -0,0 +1,187 @@
+//===- 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/Analysis/AliasAnalysis.h"
+#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 the RHS reads that same
+    // storage (e.g. `a = a(:n)` or `a = a(:n) + 1`), the freed pointer would
+    // be read by the subsequent non-realloc assign, causing use-after-free.
+    //
+    // Reading allocatable storage always goes through a `fir.load` of the
+    // descriptor, so walking the RHS def-chain and querying alias analysis
+    // on any encountered load's memref soundly detects aliasing without
+    // running into the descriptor-vs-data NoAlias quirk of fir::AliasAnalysis
----------------
jeanPerier wrote:

The assumption that you will have a fir.load is a bit too aggressive. There could very well be an impure elemental that is accessing the allocatable indirectly (which is at least possible if it is a module variable or has a TARGET attribute).

There is no way around doing a proper alias analysis. I understand that part of the difficulty is that you do not have an SSA value for the LHS data.

Another example where what you are doing is incorrect is if you have `a = returns_func_pointer()` where the function pointer returned by the function points to the allocatable.

Because the function pointer result descriptor is a local alloca that is distinct from `a` descriptor, you will get false negative on the fact that the RHS is not using the data. In other words, the data pointed by `a` may be accessed via another descriptor in the `RHS` that does not alias with `a` descriptor. So you cannot reduce the problem to a descriptor aliasing problem.

The easiest solution I see to this problem is to generate a load of the LHS fir.box and to use the alias analysis framework with that, and to erase that load afterward. The long term solution would probably to be able to add an entry point to the alias analysis that we care about the data being pointed to by a descriptor address, but this is likely a bit more involving change.

 

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