From c203850ad55de6e1396d5735e4d9b56b66db9220 Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier@nvidia.com>
Date: Tue, 25 Apr 2023 09:03:55 +0200
Subject: [PATCH] [flang][hlfir] Support fir.declare in AbstractResult pass

The AbstractResult pass replaces allocation of function result on the callee
side per an extra argument so that the allocation of the result can be
done on the caller stack.
It looks for the result allocation from the fir.return op, so it needs
to handle (in a transparent way) a fir.declare in the chain between the
allocation and the fir.return.

Reviewed By: vzakhari, clementval

Differential Revision: https://reviews.llvm.org/D149057
---
 .../Optimizer/Transforms/AbstractResult.cpp   |  6 ++++-
 flang/test/Fir/abstract-results.fir           | 27 +++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/flang/lib/Optimizer/Transforms/AbstractResult.cpp b/flang/lib/Optimizer/Transforms/AbstractResult.cpp
index 2c0576eaa5cc..7fcdb4d5650b 100644
--- a/flang/lib/Optimizer/Transforms/AbstractResult.cpp
+++ b/flang/lib/Optimizer/Transforms/AbstractResult.cpp
@@ -217,6 +217,10 @@ public:
     if (auto *op = returnedValue.getDefiningOp())
       if (auto load = mlir::dyn_cast<fir::LoadOp>(op)) {
         auto resultStorage = load.getMemref();
+        // The result alloca may be behind a fir.declare, if any.
+        if (auto declare = mlir::dyn_cast_or_null<fir::DeclareOp>(
+                resultStorage.getDefiningOp()))
+          resultStorage = declare.getMemref();
         // TODO: This should be generalized for derived types, and it is
         // architecture and OS dependent.
         if (fir::isa_builtin_cptr_type(returnedValue.getType())) {
@@ -232,7 +236,7 @@ public:
               ret, mlir::ValueRange{retValue});
           return mlir::success();
         }
-        load.getMemref().replaceAllUsesWith(newArg);
+        resultStorage.replaceAllUsesWith(newArg);
         replacedStorage = true;
         if (auto *alloc = resultStorage.getDefiningOp())
           if (alloc->use_empty())
diff --git a/flang/test/Fir/abstract-results.fir b/flang/test/Fir/abstract-results.fir
index 374c0d18753b..42ff2a5c8eb2 100644
--- a/flang/test/Fir/abstract-results.fir
+++ b/flang/test/Fir/abstract-results.fir
@@ -106,6 +106,33 @@ func.func @retcptr() -> !fir.type<_QM__fortran_builtinsT__builtin_c_ptr{__addres
   // FUNC-BOX: return %[[VAL]] : i64
 }
 
+// FUNC-REF-LABEL:  func private @arrayfunc_callee_declare(
+// FUNC-REF-SAME: %[[buffer:.*]]: !fir.ref<!fir.array<?xf32>>, %[[n:.*]]: index) {
+// FUNC-BOX-LABEL:  func private @arrayfunc_callee_declare(
+// FUNC-BOX-SAME: %[[box:.*]]: !fir.box<!fir.array<?xf32>>, %[[n:.*]]: index) {
+func.func private @arrayfunc_callee_declare(%n : index) -> !fir.array<?xf32> {
+  %buffer_alloc = fir.alloca !fir.array<?xf32>, %n
+  %shape = fir.shape %n : (index) -> !fir.shape<1>
+  %buffer = fir.declare %buffer_alloc(%shape) {uniq_name = "x"}: (!fir.ref<!fir.array<?xf32>>, !fir.shape<1>) -> !fir.ref<!fir.array<?xf32>>
+  // Do something with result (res(4) = 42.)
+  %c4 = arith.constant 4 : i64
+  %coor = fir.coordinate_of %buffer, %c4 : (!fir.ref<!fir.array<?xf32>>, i64) -> !fir.ref<f32>
+  %cst = arith.constant 4.200000e+01 : f32
+  fir.store %cst to %coor : !fir.ref<f32>
+  %res = fir.load %buffer : !fir.ref<!fir.array<?xf32>>
+  return %res : !fir.array<?xf32>
+
+  // FUNC-REF-DAG: %[[buffer_declare:.*]] = fir.declare %[[buffer]](%{{.*}}) {uniq_name = "x"} : (!fir.ref<!fir.array<?xf32>>, !fir.shape<1>) -> !fir.ref<!fir.array<?xf32>>
+  // FUNC-REF-DAG: %[[coor:.*]] = fir.coordinate_of %[[buffer_declare]], %{{.*}} : (!fir.ref<!fir.array<?xf32>>, i64) -> !fir.ref<f32>
+  // FUNC-REF-DAG: fir.store %{{.*}} to %[[coor]] : !fir.ref<f32>
+  // FUNC-REF: return
+
+  // FUNC-BOX: %[[buffer:.*]] = fir.box_addr %[[box]] : (!fir.box<!fir.array<?xf32>>) -> !fir.ref<!fir.array<?xf32>>
+  // FUNC-BOX-DAG: %[[buffer_declare:.*]] = fir.declare %[[buffer]](%{{.*}}) {uniq_name = "x"} : (!fir.ref<!fir.array<?xf32>>, !fir.shape<1>) -> !fir.ref<!fir.array<?xf32>>
+  // FUNC-BOX-DAG: %[[coor:.*]] = fir.coordinate_of %[[buffer_declare]], %{{.*}} : (!fir.ref<!fir.array<?xf32>>, i64) -> !fir.ref<f32>
+  // FUNC-BOX-DAG: fir.store %{{.*}} to %[[coor]] : !fir.ref<f32>
+  // FUNC-BOX: return
+}
 
 // ------------------------ Test caller rewrite --------------------------------
 
-- 
2.34.1