[flang] Treat hlfir.associate as Allocate for FIR alias analysis.

[vzakhari]

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.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

Author: Slava Zakharin (vzakhari)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/139004.diff

2 Files Affected:

• (modified) flang/lib/Optimizer/Analysis/AliasAnalysis.cpp (+14)
• (added) flang/test/HLFIR/opt-bufferization-eval_in_mem-with-associate.fir (+30)

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

[tblah]

LGTM.

Maybe this assumption should be documented for future reference in the hlfir.associate operation definition, because otherwise the connection between the assumption here and the implementation of HLFIR bufferization is not that obvious.

[jeanPerier]

LGTM

[vzakhari]

Thanks for the review!

Tom, I added a comment in the op definition.