[flang][hlfir] Support fir.declare in AbstractResult pass

[jeanPerier]

Summary

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.

[vzakhari]

Thank you!

[clementval]

LGTM

flang/lib/Optimizer/Transforms/AbstractResult.cpp

matchAndRewrite(mlir::func::ReturnOp ret,
                   mlir::PatternRewriter &rewriter) const override {
     auto loc = ret.getLoc();
     rewriter.setInsertionPoint(ret);
     auto returnedValue = ret.getOperand(0);
     bool replacedStorage = false;
     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())) {
           rewriter.eraseOp(load);
           auto module = ret->getParentOfType<mlir::ModuleOp>();
           fir::KindMapping kindMap = fir::getKindMapping(module);
           FirOpBuilder builder(rewriter, kindMap);
           mlir::Value retAddr = fir::factory::genCPtrOrCFunptrAddr(
               builder, loc, resultStorage, returnedValue.getType());
           mlir::Value retValue = rewriter.create<fir::LoadOp>(
               loc, fir::unwrapRefType(retAddr.getType()), retAddr);
           rewriter.replaceOpWithNewOp<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())
             rewriter.eraseOp(alloc);
       }
     // The result storage may have been optimized out by a memory to
     // register pass, this is possible for fir.box results, or fir.record
     // with no length parameters. Simply store the result in the result storage.

flang/test/Fir/abstract-results.fir

func.func @retcptr() -> !fir.type<_QM__fortran_builtinsT__builtin_c_ptr{__address:i64}> {
   // FUNC-REF: return %[[VAL]] : i64
   // FUNC-BOX: %[[ALLOC:.*]] = fir.alloca !fir.type<_QM__fortran_builtinsT__builtin_c_ptr{__address:i64}> {bindc_name = "rec", uniq_name = "_QFrecErec"}
   // FUNC-BOX: %[[FIELD:.*]] = fir.field_index __address, !fir.type<_QM__fortran_builtinsT__builtin_c_ptr{__address:i64}>
   // FUNC-BOX: %[[ADDR:.*]] = fir.coordinate_of %[[ALLOC]], %[[FIELD]] : (!fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_c_ptr{__address:i64}>>, !fir.field) -> !fir.ref<i64>
   // FUNC-BOX: %[[VAL:.*]] = fir.load %[[ADDR]] : !fir.ref<i64>
   // 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 --------------------------------
 
 // FUNC-REF-LABEL: func @call_arrayfunc() {
 // FUNC-BOX-LABEL: func @call_arrayfunc() {
 func.func @call_arrayfunc() {
   %c100 = arith.constant 100 : index
   %buffer = fir.alloca !fir.array<?xf32>, %c100