[AArch64] Map G_LOAD on FPR when the definition goes to a copy to FPR

We used to detect loads feeding fp instructions, but we were
failing to take into account cases where this happens through copies.
For instance, loads can fed copies coming from the ABI lowering
of floating point arguments/results.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@318589 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 16 insertions, 4 deletions

diff --git a/lib/Target/AArch64/AArch64RegisterBankInfo.cpp b/lib/Target/AArch64/AArch64RegisterBankInfo.cpp
index 9940a74d7a8..c497669f937 100644
--- a/lib/Target/AArch64/AArch64RegisterBankInfo.cpp
+++ b/lib/Target/AArch64/AArch64RegisterBankInfo.cpp
@@ -594,15 +594,24 @@ AArch64RegisterBankInfo::getInstrMapping(const MachineInstr &MI) const {
       // In that case, we want the default mapping to be on FPR
       // instead of blind map every scalar to GPR.
       for (const MachineInstr &UseMI :
-           MRI.use_instructions(MI.getOperand(0).getReg()))
+           MRI.use_instructions(MI.getOperand(0).getReg())) {
         // If we have at least one direct use in a FP instruction,
         // assume this was a floating point load in the IR.
         // If it was not, we would have had a bitcast before
         // reaching that instruction.
-        if (isPreISelGenericFloatingPointOpcode(UseMI.getOpcode())) {
+        unsigned UseOpc = UseMI.getOpcode();
+        if (isPreISelGenericFloatingPointOpcode(UseOpc) ||
+            // Check if we feed a copy-like instruction with
+            // floating point constraints. In that case, we are still
+            // feeding fp instructions, but indirectly
+            // (e.g., through ABI copies).
+            ((UseOpc == TargetOpcode::COPY || UseMI.isPHI()) &&
+             getRegBank(UseMI.getOperand(0).getReg(), MRI, TRI) ==
+                 &AArch64::FPRRegBank)) {
           OpRegBankIdx[0] = PMI_FirstFPR;
           break;
         }
+      }
     break;
   case TargetOpcode::G_STORE:
     // Check if that store is fed by fp instructions.
diff --git a/test/CodeGen/AArch64/GlobalISel/arm64-regbankselect.mir b/test/CodeGen/AArch64/GlobalISel/arm64-regbankselect.mir
index aace553af2d..a27cf2bea78 100644
--- a/test/CodeGen/AArch64/GlobalISel/arm64-regbankselect.mir
+++ b/test/CodeGen/AArch64/GlobalISel/arm64-regbankselect.mir
@@ -917,7 +917,7 @@ body:             |
 # CHECK: registers:
 # CHECK:  - { id: 0, class: fpr, preferred-register: '' }
 # CHECK:  - { id: 1, class: gpr, preferred-register: '' }
-# CHECK:  - { id: 2, class: gpr, preferred-register: '' }
+# CHECK:  - { id: 2, class: fpr, preferred-register: '' }
 #
 # CHECK:  %0:fpr(s16) = COPY %h0
 # CHECK-NEXT: %1:gpr(p0) = G_FRAME_INDEX %stack.0.p.addr
@@ -925,7 +925,10 @@ body:             |
 # would have been on GPR and we would have to insert a copy to move
 # the value away from FPR (h0).
 # CHECK-NEXT: G_STORE %0(s16), %1(p0) :: (store 2 into %ir.p.addr)
-# CHECK-NEXT: %2:gpr(s16) = G_LOAD %1(p0) :: (load 2 from %ir.p.addr)
+# If we didn't look through the copy for %2, the default mapping
+# would have been on GPR and we would have to insert a copy to move
+# the value to FPR (h0).
+# CHECK-NEXT: %2:fpr(s16) = G_LOAD %1(p0) :: (load 2 from %ir.p.addr)
 # CHECK-NEXT: %h0 = COPY %2(s16)
 name:            passFp16ViaAllocas
 alignment:       2