[llvm] [MachineLoopInfo] Implement `isLoopInvariant` recursively. (PR #95285)

[Mikhail Gudim via llvm-commits]

https://github.com/mgudim created https://github.com/llvm/llvm-project/pull/95285

In order for the instruction to be loop invariant, each of its operands has to be defined outside of the loop or be loop invariant itself.

>From d15dec103ae1653139bd3b6896a016b9eb6ad50f Mon Sep 17 00:00:00 2001
From: Mikhail Gudim <mgudim at gmail.com>
Date: Fri, 31 May 2024 13:56:28 -0400
Subject: [PATCH] [MachineLoopInfo] Implement `isLoopInvariant` recursively.

In order for the instruction to be loop invariant, each of its operands
has to be defined outside of the loop or be loop invariant itself.
---
 llvm/include/llvm/CodeGen/MachineLoopInfo.h |  6 ++--
 llvm/lib/CodeGen/MachineLoopInfo.cpp        | 32 +++++++++++++--------
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/MachineLoopInfo.h b/llvm/include/llvm/CodeGen/MachineLoopInfo.h
index 967c4a70ca469..dcc9f699575ad 100644
--- a/llvm/include/llvm/CodeGen/MachineLoopInfo.h
+++ b/llvm/include/llvm/CodeGen/MachineLoopInfo.h
@@ -82,7 +82,8 @@ class MachineLoop : public LoopBase<MachineBasicBlock, MachineLoop> {
   /// ExcludeReg can be used to exclude the given register from the check
   /// i.e. when we're considering hoisting it's definition but not hoisted it
   /// yet
-  bool isLoopInvariant(MachineInstr &I, const Register ExcludeReg = 0) const;
+  bool isLoopInvariant(const MachineInstr &I, const Register ExcludeReg = 0,
+                       unsigned RecursionDepth = 1) const;
 
   void dump() const;
 
@@ -90,7 +91,8 @@ class MachineLoop : public LoopBase<MachineBasicBlock, MachineLoop> {
   friend class LoopInfoBase<MachineBasicBlock, MachineLoop>;
 
   /// Returns true if the given physreg has no defs inside the loop.
-  bool isLoopInvariantImplicitPhysReg(Register Reg) const;
+  bool isLoopInvariantImplicitPhysReg(Register Reg, Register ExcludeReg,
+                                      unsigned RecursionDepth = 0) const;
 
   explicit MachineLoop(MachineBasicBlock *MBB)
     : LoopBase<MachineBasicBlock, MachineLoop>(MBB) {}
diff --git a/llvm/lib/CodeGen/MachineLoopInfo.cpp b/llvm/lib/CodeGen/MachineLoopInfo.cpp
index 1019c53e57c6f..4f8245a35c9ab 100644
--- a/llvm/lib/CodeGen/MachineLoopInfo.cpp
+++ b/llvm/lib/CodeGen/MachineLoopInfo.cpp
@@ -198,7 +198,8 @@ MDNode *MachineLoop::getLoopID() const {
   return LoopID;
 }
 
-bool MachineLoop::isLoopInvariantImplicitPhysReg(Register Reg) const {
+bool MachineLoop::isLoopInvariantImplicitPhysReg(
+    Register Reg, Register ExcludeReg, unsigned RecursionDepth) const {
   MachineFunction *MF = getHeader()->getParent();
   MachineRegisterInfo *MRI = &MF->getRegInfo();
 
@@ -210,15 +211,20 @@ bool MachineLoop::isLoopInvariantImplicitPhysReg(Register Reg) const {
            ->shouldAnalyzePhysregInMachineLoopInfo(Reg))
     return false;
 
-  return !llvm::any_of(
-      MRI->def_instructions(Reg),
-      [this](const MachineInstr &MI) { return this->contains(&MI); });
+  return !llvm::any_of(MRI->def_instructions(Reg), [=](const MachineInstr &MI) {
+    return (this->contains(&MI) &&
+            !isLoopInvariant(MI, ExcludeReg, RecursionDepth - 1));
+  });
 }
 
-bool MachineLoop::isLoopInvariant(MachineInstr &I,
-                                  const Register ExcludeReg) const {
-  MachineFunction *MF = I.getParent()->getParent();
-  MachineRegisterInfo *MRI = &MF->getRegInfo();
+bool MachineLoop::isLoopInvariant(const MachineInstr &I,
+                                  const Register ExcludeReg,
+                                  unsigned RecursionDepth) const {
+  if (RecursionDepth == 0)
+    return false;
+
+  const MachineFunction *MF = I.getParent()->getParent();
+  const MachineRegisterInfo *MRI = &MF->getRegInfo();
   const TargetSubtargetInfo &ST = MF->getSubtarget();
   const TargetRegisterInfo *TRI = ST.getRegisterInfo();
   const TargetInstrInfo *TII = ST.getInstrInfo();
@@ -243,7 +249,7 @@ bool MachineLoop::isLoopInvariant(MachineInstr &I,
         // it could get allocated to something with a def during allocation.
         // However, if the physreg is known to always be caller saved/restored
         // then this use is safe to hoist.
-        if (!isLoopInvariantImplicitPhysReg(Reg) &&
+        if (!isLoopInvariantImplicitPhysReg(Reg, ExcludeReg, RecursionDepth) &&
             !(TRI->isCallerPreservedPhysReg(Reg.asMCReg(), *I.getMF())) &&
             !TII->isIgnorableUse(MO))
           return false;
@@ -265,9 +271,11 @@ bool MachineLoop::isLoopInvariant(MachineInstr &I,
     assert(MRI->getVRegDef(Reg) &&
            "Machine instr not mapped for this vreg?!");
 
-    // If the loop contains the definition of an operand, then the instruction
-    // isn't loop invariant.
-    if (contains(MRI->getVRegDef(Reg)))
+    // If the loop contains the definition of an operand, then it must be loop
+    // invariant
+    MachineInstr *VRegDefMI = MRI->getVRegDef(Reg);
+    if (contains(VRegDefMI) &&
+        !isLoopInvariant(*VRegDefMI, ExcludeReg, RecursionDepth - 1))
       return false;
   }