From 4d5a9c1d171b1b86deb2030e37b31ab2ec9d1dd0 Mon Sep 17 00:00:00 2001
From: Renaud-K <rkauffmann@nvidia.com>
Date: Wed, 9 Nov 2022 11:04:22 -0800
Subject: [PATCH] Using higher level interface to insert new arguments so
 arguments and their corresponding attributes are moved together

Differential revision: https://reviews.llvm.org/D13767
---
 flang/lib/Optimizer/Transforms/AbstractResult.cpp |  9 ++++++--
 flang/test/Fir/abstract-result-2.fir              | 27 +++++++++++++++++++++++
 2 files changed, 34 insertions(+), 2 deletions(-)
 create mode 100644 flang/test/Fir/abstract-result-2.fir

diff --git a/flang/lib/Optimizer/Transforms/AbstractResult.cpp b/flang/lib/Optimizer/Transforms/AbstractResult.cpp
index 6cdc445..b295885 100644
--- a/flang/lib/Optimizer/Transforms/AbstractResult.cpp
+++ b/flang/lib/Optimizer/Transforms/AbstractResult.cpp
@@ -334,13 +334,14 @@ public:
             });
         return;
       }
-      func.setType(getNewFunctionType(funcTy, shouldBoxResult));
       if (!func.empty()) {
         // Insert new argument.
         mlir::OpBuilder rewriter(context);
         auto resultType = funcTy.getResult(0);
         auto argTy = getResultArgumentType(resultType, shouldBoxResult);
-        mlir::Value newArg = func.front().insertArgument(0u, argTy, loc);
+        func.insertArgument(0u, argTy, {}, loc);
+        func.eraseResult(0u);
+        mlir::Value newArg = func.getArgument(0u);
         if (mustEmboxResult(resultType, shouldBoxResult)) {
           auto bufferType = fir::ReferenceType::get(resultType);
           rewriter.setInsertionPointToStart(&func.front());
@@ -349,6 +350,10 @@ public:
         patterns.insert<ReturnOpConversion>(context, newArg);
         target.addDynamicallyLegalOp<mlir::func::ReturnOp>(
             [](mlir::func::ReturnOp ret) { return ret.operands().empty(); });
+        assert(func.getFunctionType() ==
+               getNewFunctionType(funcTy, shouldBoxResult));
+      } else {
+        func.setType(getNewFunctionType(funcTy, shouldBoxResult));
       }
     }
   }
diff --git a/flang/test/Fir/abstract-result-2.fir b/flang/test/Fir/abstract-result-2.fir
new file mode 100644
index 0000000..3c13943
--- /dev/null
+++ b/flang/test/Fir/abstract-result-2.fir
@@ -0,0 +1,27 @@
+// RUN: fir-opt %s --abstract-result-on-func-opt | FileCheck %s 
+
+// Check that the attributes are shifted along with their corresponding arguments
+
+//CHECK: func.func @_QMi8Pintrinsic_pack0(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>, %arg1: !fir.box<!fir.array<?x?x?xi32>> {fir.bindc_name = "array"}, %arg2: !fir.ref<!fir.logical<4>> {fir.bindc_name = "mask"}, %arg3: !fir.box<!fir.array<?xi32>> {fir.bindc_name = "vector", fir.optional})
+
+func.func @_QMi8Pintrinsic_pack0(%arg0: !fir.box<!fir.array<?x?x?xi32>> {fir.bindc_name = "array"}, %arg1: !fir.ref<!fir.logical<4>> {fir.bindc_name = "mask"}, %arg2: !fir.box<!fir.array<?xi32>> {fir.bindc_name = "vector", fir.optional}) -> !fir.box<!fir.heap<!fir.array<?xi32>>> {
+  %0 = fir.alloca !fir.box<!fir.heap<!fir.array<?xi32>>>
+  %1 = fir.alloca !fir.box<!fir.heap<!fir.array<?xi32>>> {bindc_name = "intrinsic_pack0", uniq_name = "_QMi8Fintrinsic_pack0Eintrinsic_pack0"}
+  %2 = fir.zero_bits !fir.heap<!fir.array<?xi32>>
+  %c0 = arith.constant 0 : index
+  %3 = fir.shape %c0 : (index) -> !fir.shape<1>
+  %4 = fir.embox %2(%3) : (!fir.heap<!fir.array<?xi32>>, !fir.shape<1>) -> !fir.box<!fir.heap<!fir.array<?xi32>>>
+  fir.store %4 to %1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+  %5 = fir.embox %arg1 : (!fir.ref<!fir.logical<4>>) -> !fir.box<!fir.logical<4>>
+  %6 = fir.zero_bits !fir.heap<!fir.array<?xi32>>
+  %7 = fir.embox %6(%3) : (!fir.heap<!fir.array<?xi32>>, !fir.shape<1>) -> !fir.box<!fir.heap<!fir.array<?xi32>>>
+  fir.store %7 to %0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+  %8 = fir.convert %0 : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>) -> !fir.ref<!fir.box<none>>
+  %9 = fir.convert %arg0 : (!fir.box<!fir.array<?x?x?xi32>>) -> !fir.box<none>
+  %10 = fir.convert %5 : (!fir.box<!fir.logical<4>>) -> !fir.box<none>
+  %11 = fir.convert %arg2 : (!fir.box<!fir.array<?xi32>>) -> !fir.box<none>
+  %12 = fir.call @_FortranAPack(%8, %9, %10, %11) : (!fir.ref<!fir.box<none>>, !fir.box<none>, !fir.box<none>, !fir.box<none>) -> none
+  %13 = fir.load %1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+  return %13 : !fir.box<!fir.heap<!fir.array<?xi32>>>
+}
+func.func private @_FortranAPack(!fir.ref<!fir.box<none>>, !fir.box<none>, !fir.box<none>, !fir.box<none>) -> none attributes {fir.runtime}
-- 
2.7.4