[flang-commits] [flang] [flang] Treat hlfir.associate as Allocate for FIR alias analysis. (PR #139004)

[Slava Zakharin via flang-commits]

https://github.com/vzakhari updated https://github.com/llvm/llvm-project/pull/139004

>From ae126b8ec224d6ae963ee666839cdef70fd5e4ab Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin at nvidia.com>
Date: Wed, 7 May 2025 18:36:10 -0700
Subject: [PATCH 1/2] [flang] Treat hlfir.associate as Allocate for FIR alias
 analysis.

Early HLFIR optimizations may experience problems with values
produced by hlfir.associate. In most cases this is a unique
local memory allocation, but it can also reuse some other
hlfir.expr memory sometimes. It seems to be safe to assume
unique allocation for trivial types, since we always
allocate new memory for them.
---
 .../lib/Optimizer/Analysis/AliasAnalysis.cpp  | 14 +++++++++
 ...fferization-eval_in_mem-with-associate.fir | 30 +++++++++++++++++++
 2 files changed, 44 insertions(+)
 create mode 100644 flang/test/HLFIR/opt-bufferization-eval_in_mem-with-associate.fir

diff --git a/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp b/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp
index cbfc8b63ab64d..73ddd1ff80126 100644
--- a/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp
+++ b/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp
@@ -540,6 +540,20 @@ AliasAnalysis::Source AliasAnalysis::getSource(mlir::Value v,
           v = op.getVar();
           defOp = v.getDefiningOp();
         })
+        .Case<hlfir::AssociateOp>([&](auto op) {
+          mlir::Value source = op.getSource();
+          if (fir::isa_trivial(source.getType())) {
+            // Trivial values will always use distinct temp memory,
+            // so we can classify this as Allocate and stop.
+            type = SourceKind::Allocate;
+            breakFromLoop = true;
+          } else {
+            // AssociateOp may reuse the expression storage,
+            // so we have to trace further.
+            v = source;
+            defOp = v.getDefiningOp();
+          }
+        })
         .Case<fir::AllocaOp, fir::AllocMemOp>([&](auto op) {
           // Unique memory allocation.
           type = SourceKind::Allocate;
diff --git a/flang/test/HLFIR/opt-bufferization-eval_in_mem-with-associate.fir b/flang/test/HLFIR/opt-bufferization-eval_in_mem-with-associate.fir
new file mode 100644
index 0000000000000..2f5f88ff9e7e8
--- /dev/null
+++ b/flang/test/HLFIR/opt-bufferization-eval_in_mem-with-associate.fir
@@ -0,0 +1,30 @@
+// RUN: fir-opt --opt-bufferization %s | FileCheck %s
+
+// Verify that hlfir.eval_in_mem uses the LHS array instead
+// of allocating a temporary.
+func.func @_QPtest() {
+  %cst = arith.constant 1.000000e+00 : f32
+  %c10 = arith.constant 10 : index
+  %0 = fir.dummy_scope : !fir.dscope
+  %1 = fir.alloca !fir.array<10xf32> {bindc_name = "x", uniq_name = "_QFtestEx"}
+  %2 = fir.shape %c10 : (index) -> !fir.shape<1>
+  %3:2 = hlfir.declare %1(%2) {uniq_name = "_QFtestEx"} : (!fir.ref<!fir.array<10xf32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xf32>>, !fir.ref<!fir.array<10xf32>>)
+  %4:3 = hlfir.associate %cst {adapt.valuebyref} : (f32) -> (!fir.ref<f32>, !fir.ref<f32>, i1)
+  %5 = hlfir.eval_in_mem shape %2 : (!fir.shape<1>) -> !hlfir.expr<10xf32> {
+  ^bb0(%arg0: !fir.ref<!fir.array<10xf32>>):
+    %6 = fir.call @_QParray_func(%4#0) fastmath<contract> : (!fir.ref<f32>) -> !fir.array<10xf32>
+    fir.save_result %6 to %arg0(%2) : !fir.array<10xf32>, !fir.ref<!fir.array<10xf32>>, !fir.shape<1>
+  }
+  hlfir.assign %5 to %3#0 : !hlfir.expr<10xf32>, !fir.ref<!fir.array<10xf32>>
+  hlfir.end_associate %4#1, %4#2 : !fir.ref<f32>, i1
+  hlfir.destroy %5 : !hlfir.expr<10xf32>
+  return
+}
+// CHECK-LABEL:   func.func @_QPtest() {
+// CHECK:           %[[VAL_0:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK:           %[[VAL_3:.*]] = fir.alloca !fir.array<10xf32> {bindc_name = "x", uniq_name = "_QFtestEx"}
+// CHECK:           %[[VAL_5:.*]]:2 = hlfir.declare %[[VAL_3]](%[[VAL_4:.*]]) {uniq_name = "_QFtestEx"} : (!fir.ref<!fir.array<10xf32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xf32>>, !fir.ref<!fir.array<10xf32>>)
+// CHECK:           %[[VAL_6:.*]]:3 = hlfir.associate %[[VAL_0]] {adapt.valuebyref} : (f32) -> (!fir.ref<f32>, !fir.ref<f32>, i1)
+// CHECK:           %[[VAL_7:.*]] = fir.call @_QParray_func(%[[VAL_6]]#0) fastmath<contract> : (!fir.ref<f32>) -> !fir.array<10xf32>
+// CHECK:           fir.save_result %[[VAL_7]] to %[[VAL_5]]#0(%[[VAL_4]]) : !fir.array<10xf32>, !fir.ref<!fir.array<10xf32>>, !fir.shape<1>
+// CHECK:           hlfir.end_associate %[[VAL_6]]#1, %[[VAL_6]]#2 : !fir.ref<f32>, i1

>From 451712c10e80d987091e3411dfe42fd6542d19f7 Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin at nvidia.com>
Date: Fri, 9 May 2025 15:06:32 -0700
Subject: [PATCH 2/2] Added aliasing comment in hlfir.associate definition.

---
 flang/include/flang/Optimizer/HLFIR/HLFIROps.td | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
index f69930d5b53b3..08cfa8168a922 100644
--- a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
+++ b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
@@ -758,6 +758,13 @@ def hlfir_AssociateOp : hlfir_Op<"associate", [AttrSizedOperandSegments,
     For expressions, this operation is an incentive to re-use the expression
     storage, if any, after the bufferization pass when possible (if the
     expression is not used afterwards).
+
+    For aliasing purposes, hlfir.associate with the source being
+    a trivial FIR value is considered to be a unique allocation
+    that does not alias with anything else. For non-trivial cases
+    it may be a unique allocation or an alias for the source expression
+    storage, so FIR alias analysis will look through it for finding
+    the source.
   }];
 
   let arguments = (ins