[flang-commits] [PATCH] D149057: [flang][hlfir] Support fir.declare in AbstractResult pass

[Jean Perier via Phabricator via flang-commits]

jeanPerier created this revision.
jeanPerier added reviewers: vzakhari, clementval.
jeanPerier added a project: Flang.
Herald added subscribers: sunshaoce, mehdi_amini, jdoerfert.
Herald added a project: All.
jeanPerier requested review of this revision.

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.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D149057

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


Index: flang/test/Fir/abstract-results.fir
===================================================================
--- flang/test/Fir/abstract-results.fir
+++ flang/test/Fir/abstract-results.fir
@@ -106,6 +106,33 @@
   // 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 --------------------------------
 
Index: flang/lib/Optimizer/Transforms/AbstractResult.cpp
===================================================================
--- flang/lib/Optimizer/Transforms/AbstractResult.cpp
+++ flang/lib/Optimizer/Transforms/AbstractResult.cpp
@@ -217,6 +217,10 @@
     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 @@
               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())


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D149057.516361.patch
Type: text/x-patch
Size: 3406 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/flang-commits/attachments/20230424/45d4d682/attachment-0001.bin>