[llvm] [MachineLICM] Hoist COPY instruction only when user can be hoisted (PR #81735)

[via llvm-commits]

https://github.com/michaelselehov created https://github.com/llvm/llvm-project/pull/81735

befa925acac8fd6a9266e introduced preliminary hoisting of COPY instructions when the user of the COPY is inside the same loop. That optimization appeared to be too aggressive and hoisted too many COPY's greatly increasing register pressure causing performance regressions for AMDGPU target.

This is intended to fix the regression by hoisting COPY instruction only if either:
 - User of COPY can be hoisted (other args are invariant) 
 or
 - Hoisting COPY doesn't bring high register pressure

>From 65f08edbfae23519ab220d11d6748b158fd8788d Mon Sep 17 00:00:00 2001
From: Michael Selehov <michael.selehov at amd.com>
Date: Fri, 26 Jan 2024 16:58:34 +0000
Subject: [PATCH] [MachineLICM] Hoist COPY instruction only when user can be
 hoisted

befa925acac8fd6a9266e introduced preliminary hoisting of COPY
instructions when the user of the COPY is inside the same loop.
That optimization appeared to be too aggressive and hoisted too many
COPY's greately increasing register pressure causing performance
regressions for AMDGPU target.

This is intended to fix the regression by hoisting COPY instruction only
if either:
 - User of COPY can be hoisted (other args are invariant)
or
 - Hoisting COPY doesn't bring high register pressure
---
 llvm/include/llvm/CodeGen/MachineLoopInfo.h |  4 +++-
 llvm/lib/CodeGen/MachineLICM.cpp            | 16 +++++++++++++---
 llvm/lib/CodeGen/MachineLoopInfo.cpp        |  4 +++-
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/MachineLoopInfo.h b/llvm/include/llvm/CodeGen/MachineLoopInfo.h
index ae075bee1daf76..f3811cd65a3167 100644
--- a/llvm/include/llvm/CodeGen/MachineLoopInfo.h
+++ b/llvm/include/llvm/CodeGen/MachineLoopInfo.h
@@ -79,7 +79,9 @@ class MachineLoop : public LoopBase<MachineBasicBlock, MachineLoop> {
   /// I.e., all virtual register operands are defined outside of the loop,
   /// physical registers aren't accessed explicitly, and there are no side
   /// effects that aren't captured by the operands or other flags.
-  bool isLoopInvariant(MachineInstr &I) const;
+  /// 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;
 
   void dump() const;
 
diff --git a/llvm/lib/CodeGen/MachineLICM.cpp b/llvm/lib/CodeGen/MachineLICM.cpp
index efc19f8fdbf8c2..e9cbf4e8c8297f 100644
--- a/llvm/lib/CodeGen/MachineLICM.cpp
+++ b/llvm/lib/CodeGen/MachineLICM.cpp
@@ -1264,13 +1264,23 @@ bool MachineLICMBase::IsProfitableToHoist(MachineInstr &MI,
 
   // If we have a COPY with other uses in the loop, hoist to allow the users to
   // also be hoisted.
+  Register defReg;
   if (MI.isCopy() && MI.getOperand(0).isReg() &&
-      MI.getOperand(0).getReg().isVirtual() && MI.getOperand(1).isReg() &&
+      (defReg = MI.getOperand(0).getReg()).isVirtual() && MI.getOperand(1).isReg() &&
       MI.getOperand(1).getReg().isVirtual() &&
       IsLoopInvariantInst(MI, CurLoop) &&
       any_of(MRI->use_nodbg_instructions(MI.getOperand(0).getReg()),
-             [&CurLoop](MachineInstr &UseMI) {
-               return CurLoop->contains(&UseMI);
+            [&CurLoop, this, defReg, Cost](MachineInstr &UseMI) {
+                  if (!CurLoop->contains(&UseMI))
+                    return false;
+
+                  // COPY is a cheap instruction, but if moving it won't cause high RP we're
+                  // fine to hoist it even if the user can't be hoisted later
+                  // Otherwise we want to check the user if it's hoistable
+                  if (CanCauseHighRegPressure(Cost, false) && !CurLoop->isLoopInvariant(UseMI, defReg))
+                    return false;
+
+                  return true;
              }))
     return true;
 
diff --git a/llvm/lib/CodeGen/MachineLoopInfo.cpp b/llvm/lib/CodeGen/MachineLoopInfo.cpp
index bdbc57099aa8d3..0d713ba3d6e23c 100644
--- a/llvm/lib/CodeGen/MachineLoopInfo.cpp
+++ b/llvm/lib/CodeGen/MachineLoopInfo.cpp
@@ -198,7 +198,7 @@ MDNode *MachineLoop::getLoopID() const {
   return LoopID;
 }
 
-bool MachineLoop::isLoopInvariant(MachineInstr &I) const {
+bool MachineLoop::isLoopInvariant(MachineInstr &I, const Register ExcludeReg) const {
   MachineFunction *MF = I.getParent()->getParent();
   MachineRegisterInfo *MRI = &MF->getRegInfo();
   const TargetSubtargetInfo &ST = MF->getSubtarget();
@@ -213,6 +213,8 @@ bool MachineLoop::isLoopInvariant(MachineInstr &I) const {
     Register Reg = MO.getReg();
     if (Reg == 0) continue;
 
+    if (ExcludeReg == Reg) continue;
+
     // An instruction that uses or defines a physical register can't e.g. be
     // hoisted, so mark this as not invariant.
     if (Reg.isPhysical()) {