[MachineSink][AArch64] Sink instruction copies when they can replace copy into hard register or folded into addressing mode

This patch adds a new code transformation to the `MachineSink` pass,
that tries to sink copies of an instruction, when the copies can be folded
into the addressing modes of load/store instructions, or
replace another instruction (currently, copies into a hard register).

The criteria for performing the transformation is that:
* the register pressure at the sink destination block must not
  exceed the register pressure limits
* the latency and throughput of the load/store or the copy must not deteriorate
* the original instruction must be deleted

Reviewed By: dmgreen

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

Author: momchil-velikov

llvm/include/llvm/CodeGen/TargetInstrInfo.h

@@ -85,11 +85,21 @@ struct RegImmPair {
 
 /// Used to describe addressing mode similar to ExtAddrMode in CodeGenPrepare.
 /// It holds the register values, the scale value and the displacement.
+/// It also holds a descriptor for the expression used to calculate the address
+/// from the operands.
 struct ExtAddrMode {
+  enum class Formula {
+    Basic = 0,         // BaseReg + ScaledReg * Scale + Displacement
+    SExtScaledReg = 1, // BaseReg + sext(ScaledReg) * Scale + Displacement
+    ZExtScaledReg = 2  // BaseReg + zext(ScaledReg) * Scale + Displacement
+  };
+
   Register BaseReg;
   Register ScaledReg;
-  int64_t Scale;
-  int64_t Displacement;
+  int64_t Scale = 0;
+  int64_t Displacement = 0;
+  Formula Form = Formula::Basic;
+  ExtAddrMode() = default;
 };
 
 //---------------------------------------------------------------------------
@@ -1436,6 +1446,26 @@ class TargetInstrInfo : public MCInstrInfo {
     return std::nullopt;
   }
 
+  /// Check if it's possible and beneficial to fold the addressing computation
+  /// `AddrI` into the addressing mode of the load/store instruction `MemI`. The
+  /// memory instruction is a user of the virtual register `Reg`, which in turn
+  /// is the ultimate destination of zero or more COPY instructions from the
+  /// output register of `AddrI`.
+  /// Return the adddressing mode after folding in `AM`.
+  virtual bool canFoldIntoAddrMode(const MachineInstr &MemI, Register Reg,
+                                   const MachineInstr &AddrI,
+                                   ExtAddrMode &AM) const {
+    return false;
+  }
+
+  /// Emit a load/store instruction with the same value register as `MemI`, but
+  /// using the address from `AM`. The addressing mode must have been obtained
+  /// from `canFoldIntoAddr` for the same memory instruction.
+  virtual MachineInstr *emitLdStWithAddr(MachineInstr &MemI,
+                                         const ExtAddrMode &AM) const {
+    llvm_unreachable("target did not implement emitLdStWithAddr()");
+  }
+
   /// Returns true if MI's Def is NullValueReg, and the MI
   /// does not change the Zero value. i.e. cases such as rax = shr rax, X where
   /// NullValueReg = rax. Note that if the NullValueReg is non-zero, this

llvm/include/llvm/CodeGen/TargetPassConfig.h

@@ -130,6 +130,11 @@ class TargetPassConfig : public ImmutablePass {
   /// Default setting for -enable-tail-merge on this target.
   bool EnableTailMerge = true;
 
+  /// Enable sinking of instructions in MachineSink where a computation can be
+  /// folded into the addressing mode of a memory load/store instruction or
+  /// replace a copy.
+  bool EnableSinkAndFold = false;
+
   /// Require processing of functions such that callees are generated before
   /// callers.
   bool RequireCodeGenSCCOrder = false;
@@ -176,6 +181,9 @@ class TargetPassConfig : public ImmutablePass {
   bool getEnableTailMerge() const { return EnableTailMerge; }
   void setEnableTailMerge(bool Enable) { setOpt(EnableTailMerge, Enable); }
 
+  bool getEnableSinkAndFold() const { return EnableSinkAndFold; }
+  void setEnableSinkAndFold(bool Enable) { setOpt(EnableSinkAndFold, Enable); }
+
   bool requiresCodeGenSCCOrder() const { return RequireCodeGenSCCOrder; }
   void setRequiresCodeGenSCCOrder(bool Enable = true) {
     setOpt(RequireCodeGenSCCOrder, Enable);

llvm/lib/CodeGen/ImplicitNullChecks.cpp

@@ -372,7 +372,7 @@ ImplicitNullChecks::isSuitableMemoryOp(const MachineInstr &MI,
   if (!MI.mayLoadOrStore() || MI.isPredicable())
     return SR_Unsuitable;
   auto AM = TII->getAddrModeFromMemoryOp(MI, TRI);
-  if (!AM)
+  if (!AM || AM->Form != ExtAddrMode::Formula::Basic)
     return SR_Unsuitable;
   auto AddrMode = *AM;
   const Register BaseReg = AddrMode.BaseReg, ScaledReg = AddrMode.ScaledReg;