[AArch64][GISel] Assign G_LOAD defs into FPR if they feed FP instructions indirectly via PHI (#94618)

If G_LOAD def is used by a PHI, recursively check if the def of PHI is
used by any instruction that only uses FP oprand. If so, assign the def
of G_LOAD to FPR.

In other words, this change helps RBS to assign G_LOAD defs into FPR if
they are later on used in an FP instruction **indirectly via PHI**, to
avoid unnecessary copies. Before this patch, we only considered direct
usages.

It helps to fix GISel regression in TSVC kernel s3110.
GISel was 50% slower than SDAG before this change. Now GISel and SDAG
have the same runtime.

Author: Him188

llvm/lib/Target/AArch64/GISel/AArch64RegisterBankInfo.cpp

@@ -496,6 +496,20 @@ static bool isFPIntrinsic(const MachineRegisterInfo &MRI,
   }
 }
 
+bool AArch64RegisterBankInfo::isPHIWithFPContraints(
+    const MachineInstr &MI, const MachineRegisterInfo &MRI,
+    const TargetRegisterInfo &TRI, const unsigned Depth) const {
+  if (!MI.isPHI() || Depth > MaxFPRSearchDepth)
+    return false;
+
+  return any_of(MRI.use_nodbg_instructions(MI.getOperand(0).getReg()),
+                [&](const MachineInstr &UseMI) {
+                  if (onlyUsesFP(UseMI, MRI, TRI, Depth + 1))
+                    return true;
+                  return isPHIWithFPContraints(UseMI, MRI, TRI, Depth + 1);
+                });
+}
+
 bool AArch64RegisterBankInfo::hasFPConstraints(const MachineInstr &MI,
                                                const MachineRegisterInfo &MRI,
                                                const TargetRegisterInfo &TRI,
@@ -851,13 +865,18 @@ AArch64RegisterBankInfo::getInstrMapping(const MachineInstr &MI) const {
     // instead of blind map every scalar to GPR.
     if (any_of(MRI.use_nodbg_instructions(MI.getOperand(0).getReg()),
                [&](const MachineInstr &UseMI) {
-                 // If we have at least one direct use in a FP instruction,
+                 // If we have at least one direct or indirect 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.
                  //
                  // Int->FP conversion operations are also captured in
                  // onlyDefinesFP().
+
+                 if (isPHIWithFPContraints(UseMI, MRI, TRI))
+                   return true;
+
                  return onlyUsesFP(UseMI, MRI, TRI) ||
                         onlyDefinesFP(UseMI, MRI, TRI);
                }))

llvm/lib/Target/AArch64/GISel/AArch64RegisterBankInfo.h

@@ -120,6 +120,13 @@ class AArch64RegisterBankInfo final : public AArch64GenRegisterBankInfo {
   /// Maximum recursion depth for hasFPConstraints.
   const unsigned MaxFPRSearchDepth = 2;
 
+  /// \returns true if \p MI is a PHI that its def is used by
+  /// any instruction that onlyUsesFP.
+  bool isPHIWithFPContraints(const MachineInstr &MI,
+                             const MachineRegisterInfo &MRI,
+                             const TargetRegisterInfo &TRI,
+                             unsigned Depth = 0) const;
+
   /// \returns true if \p MI only uses and defines FPRs.
   bool hasFPConstraints(const MachineInstr &MI, const MachineRegisterInfo &MRI,
                      const TargetRegisterInfo &TRI, unsigned Depth = 0) const;

llvm/test/CodeGen/AArch64/GlobalISel/regbank-fp-use-def.mir

@@ -150,7 +150,7 @@ body:             |
 
 ...
 ---
-name:            load_used_by_phi_gpr
+name:            load_used_by_phi_gpr_copy_gpr
 legalized:       true
 tracksRegLiveness: true
 body:             |
@@ -173,6 +173,50 @@ body:             |
   ; CHECK-NEXT: {{  $}}
   ; CHECK-NEXT: bb.2:
   ; CHECK-NEXT:   %phi:gpr(s32) = G_PHI %gpr_copy(s32), %bb.0, %load(s32), %bb.1
+  ; CHECK-NEXT:   $w2 = COPY %phi(s32)
+  ; CHECK-NEXT:   RET_ReallyLR implicit $w2
+  bb.0:
+    successors: %bb.1(0x40000000), %bb.2(0x40000000)
+    liveins: $x0, $s0, $s1, $w0, $w1
+    %cond_wide:_(s32) = COPY $w0
+    %gpr_copy:_(s32) = COPY $w1
+    %ptr:_(p0) = COPY $x0
+    G_BRCOND %cond_wide, %bb.1
+    G_BR %bb.2
+  bb.1:
+    successors: %bb.2
+    %load:_(s32) = G_LOAD %ptr(p0) :: (load (s32))
+    G_BR %bb.2
+  bb.2:
+    %phi:_(s32) = G_PHI %gpr_copy(s32), %bb.0, %load(s32), %bb.1
+    $w2 = COPY %phi(s32)
+    RET_ReallyLR implicit $w2
+
+...
+---
+name:            load_used_by_phi_gpr_copy_fpr
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: load_used_by_phi_gpr
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
+  ; CHECK-NEXT:   liveins: $x0, $s0, $s1, $w0, $w1
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   %cond_wide:gpr(s32) = COPY $w0
+  ; CHECK-NEXT:   %gpr_copy:gpr(s32) = COPY $w1
+  ; CHECK-NEXT:   %ptr:gpr(p0) = COPY $x0
+  ; CHECK-NEXT:   G_BRCOND %cond_wide(s32), %bb.1
+  ; CHECK-NEXT:   G_BR %bb.2
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   successors: %bb.2(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   %load:fpr(s32) = G_LOAD %ptr(p0) :: (load (s32))
+  ; CHECK-NEXT:   G_BR %bb.2
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2:
+  ; CHECK-NEXT:   %phi:gpr(s32) = G_PHI %gpr_copy(s32), %bb.0, %load(s32), %bb.1
   ; CHECK-NEXT:   $s0 = COPY %phi(s32)
   ; CHECK-NEXT:   RET_ReallyLR implicit $s0
   bb.0:
@@ -189,7 +233,7 @@ body:             |
     G_BR %bb.2
   bb.2:
     %phi:_(s32) = G_PHI %gpr_copy(s32), %bb.0, %load(s32), %bb.1
-    $s0 = COPY %phi(s32)
+    $s0 = COPY %phi(s32) ; G_LOAD should consider this FPR constraint and assign %load FPR
     RET_ReallyLR implicit $s0
 
 ...