[flang-commits] [flang] c203850 - [flang][hlfir] Support fir.declare in AbstractResult pass

[Jean Perier via flang-commits]

Author: Jean Perier
Date: 2023-04-25T09:04:38+02:00
New Revision: c203850ad55de6e1396d5735e4d9b56b66db9220

URL: https://github.com/llvm/llvm-project/commit/c203850ad55de6e1396d5735e4d9b56b66db9220
DIFF: https://github.com/llvm/llvm-project/commit/c203850ad55de6e1396d5735e4d9b56b66db9220.diff

LOG: [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

Added: 
    

Modified: 
    flang/lib/Optimizer/Transforms/AbstractResult.cpp
    flang/test/Fir/abstract-results.fir

Removed: 
    


################################################################################
diff  --git a/flang/lib/Optimizer/Transforms/AbstractResult.cpp b/flang/lib/Optimizer/Transforms/AbstractResult.cpp
index 2c0576eaa5cc4..7fcdb4d5650be 100644
--- a/flang/lib/Optimizer/Transforms/AbstractResult.cpp
+++ b/flang/lib/Optimizer/Transforms/AbstractResult.cpp
@@ -217,6 +217,10 @@ class ReturnOpConversion : public mlir::OpRewritePattern<mlir::func::ReturnOp> {
     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 @@ class ReturnOpConversion : public mlir::OpRewritePattern<mlir::func::ReturnOp> {
               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 374c0d18753bb..42ff2a5c8eb2a 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 --------------------------------