[flang-commits] [flang] 5dbeb2c - [flang] Shift argument attributes when result operand is inserted

[Valentin Clement via flang-commits]

Author: Valentin Clement
Date: 2022-08-19T08:40:27+02:00
New Revision: 5dbeb2ce2a8c73016a7e1a5f0e049469239e74c7

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

LOG: [flang] Shift argument attributes when result operand is inserted

The TargetRewrite pass convert the signature of the function.
In some cases it adds operands to the function to hanlde the result of it.
This patch makes sure the argument attributes present before the conversion
are replaced with the correct arguments after the conversion is performed.

Depends D132113

Reviewed By: vdonaldson

Differential Revision: https://reviews.llvm.org/D132114

Added: 
    

Modified: 
    flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
    flang/test/Fir/target-rewrite-arg-position.fir

Removed: 
    


################################################################################
diff  --git a/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp b/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
index d87701dc51df2..b5c7374de4928 100644
--- a/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
+++ b/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
@@ -475,9 +475,20 @@ class TargetRewrite : public fir::TargetRewriteBase<TargetRewrite> {
       return;
     llvm::SmallVector<mlir::Type> newResTys;
     llvm::SmallVector<mlir::Type> newInTys;
+    llvm::SmallVector<std::pair<unsigned, mlir::NamedAttribute>> savedAttrs;
     llvm::SmallVector<std::pair<unsigned, mlir::NamedAttribute>> extraAttrs;
     llvm::SmallVector<FixupTy> fixups;
 
+    // Save argument attributes in case there is a shift so we can replace them
+    // correctly.
+    for (auto e : llvm::enumerate(funcTy.getInputs())) {
+      unsigned index = e.index();
+      llvm::ArrayRef<mlir::NamedAttribute> attrs = func.getArgAttrs(index);
+      for (mlir::NamedAttribute attr : attrs) {
+        savedAttrs.push_back({index, attr});
+      }
+    }
+
     // Convert return value(s)
     for (auto ty : funcTy.getResults())
       llvm::TypeSwitch<mlir::Type>(ty)
@@ -495,6 +506,10 @@ class TargetRewrite : public fir::TargetRewriteBase<TargetRewrite> {
           })
           .Default([&](mlir::Type ty) { newResTys.push_back(ty); });
 
+    // Saved potential shift in argument. Handling of result can add arguments
+    // at the beginning of the function signature.
+    unsigned argumentShift = newInTys.size();
+
     // Convert arguments
     llvm::SmallVector<mlir::Type> trailingTys;
     for (auto e : llvm::enumerate(funcTy.getInputs())) {
@@ -724,6 +739,17 @@ class TargetRewrite : public fir::TargetRewriteBase<TargetRewrite> {
       func.setArgAttr(extraAttr.first, extraAttr.second.getName(),
                       extraAttr.second.getValue());
 
+    // Replace attributes to the correct argument if there was an argument shift
+    // to the right.
+    if (argumentShift > 0) {
+      for (std::pair<unsigned, mlir::NamedAttribute> savedAttr : savedAttrs) {
+        func.removeArgAttr(savedAttr.first, savedAttr.second.getName());
+        func.setArgAttr(savedAttr.first + argumentShift,
+                        savedAttr.second.getName(),
+                        savedAttr.second.getValue());
+      }
+    }
+
     for (auto &fixup : fixups)
       if (fixup.finalizer)
         (*fixup.finalizer)(func);

diff  --git a/flang/test/Fir/target-rewrite-arg-position.fir b/flang/test/Fir/target-rewrite-arg-position.fir
index 1b01e73173ddf..e5359ee5540fb 100644
--- a/flang/test/Fir/target-rewrite-arg-position.fir
+++ b/flang/test/Fir/target-rewrite-arg-position.fir
@@ -1,5 +1,7 @@
 // RUN: fir-opt --target-rewrite="target=x86_64-unknown-linux-gnu" %s | FileCheck %s
 
+// Test with an argument shift.
+
 func.func @_QFPf(%arg0: !fir.ref<tuple<!fir.ref<i32>>> {fir.host_assoc}) -> !fir.complex<16> {
   %0 = fir.alloca !fir.complex<16> {bindc_name = "f", uniq_name = "_QFfEf"}
   %c2_i32 = arith.constant 2 : i32
@@ -14,4 +16,14 @@ func.func @_QFPf(%arg0: !fir.ref<tuple<!fir.ref<i32>>> {fir.host_assoc}) -> !fir
 }
 
 // CHECK-LABEL: func.func @_QFPf
-// CHECK-SAME:    %{{.*}}: !fir.ref<tuple<!fir.real<16>, !fir.real<16>>> {fir.host_assoc, llvm.align = 16 : i32, llvm.sret}, %arg1: !fir.ref<tuple<!fir.ref<i32>>> {llvm.nest}) {
+// CHECK-SAME:    %{{.*}}: !fir.ref<tuple<!fir.real<16>, !fir.real<16>>> {llvm.align = 16 : i32, llvm.sret}, %arg1: !fir.ref<tuple<!fir.ref<i32>>> {fir.host_assoc, llvm.nest}) {
+
+// -----
+
+// Test with no shift.
+
+func.func @_QFPs(%arg0: !fir.ref<i32> {fir.host_assoc}) {
+  return
+}
+
+// CHECK: func.func @_QFPs(%arg0: !fir.ref<i32> {fir.host_assoc, llvm.nest})