[CodeGenPrepare] Relax conditions for folding addressing mode into loads/stores

[chill]

Summary

The sinking of address computations to their users (loads/stores)
is often blocked by call instructions, which take the address as
a parameter - unless the call is "cold", it's considered a non-foldable use.

Considering the whole call sequence, including passing the arguments,
it is sometimes possible to materialize an address computation directly
into a hard register, in a sense "to fold the addressing mode into the call".

For example, on AArch64 the register-to-register copy
instruction ("C6.2.190 MOV (register)", which would likely by used to pass
a pre-computed address argument, is an alias to "C6.2.207 ORR (shifted register)"
and typically has the same latency and throughput as an "ADD" instruction.

This change tries to allow sinking of more addresses to load/store instructions
by preventing some call instructions from being blockers.

With this change CodeGenPrepare still does sinking only towards memory
loads/stores. It works in synergy with a MachineSink patch in
https://reviews.llvm.org/D145706, which does sinking towards calls.

This patch (together with the others up/down the stack) improves
SPECv6 500.perlbench_r by about 3.26% and the whole
of SPECv6 intrate by about 0.46% (geomean).

[efriedma]

If I'm understanding correctly, the point is that we don't want to block sinking if an address computation has multiple uses, where only some are foldable?

Expressing this in terms of "folding into a call" seems confusing and unnecessary; we should just change the logic to allow sinking if we can fold some, but not all, the uses. (Maybe including some sort analysis of whether sinking increases the number of times we perform the address computation at runtime.)

[chill]

In D143898#4123601, @efriedma wrote:

If I'm understanding correctly, the point is that we don't want to block sinking if an address computation has multiple uses, where only some are foldable?

Yes, for when we cannot say we aren't extending live ranges of the address "registers".

Expressing this in terms of "folding into a call" seems confusing and unnecessary;

To me it looks like a straightforward analogy.
For load/stores we have an addressing computation that's essentially free
when it's a part of the load/store instruction, as opposed to it being a separate instruction and using a simple indirect load/store.
For calls we have an addressing computation that's essentially free
when it's a part of the call sequence, as opposed to it being a separate instruction and using a simple register-to-register move.

we should just change the logic to allow sinking if we can fold some, but not all, the uses.
(Maybe including some sort analysis of whether sinking increases the number of times we perform the address computation at runtime.)

IMHO the main issue here is not extending live ranges too much. As for the runtime overhead, sinking a foldable address is assumed to
not increase the execution time, because of the checks in isLegalAddressingMode and (with this patch) in canFoldAddrModeIntoCall.

Perhaps, with PGO, we could use block frequencies in a way that a non-foldable, infrequently executed user does not prevent sinking, much like
the way we currently use Attribute::Cold. However, I don't think that's an alternative to this patch.

[efriedma]

For calls we have an addressing computation that's essentially free when it's a part of the call sequence, as opposed to it being a separate instruction and using a simple register-to-register move.

There are a lot of cases where there isn't any extra "mov" that can hide the cost; the most common of those being the case where the value in question is spilled. But I guess I can see the analogy. I'd generally prefer to consider that sort of transform in terms of rematerialization, though. We can compute the cost of remat much more accurately in the register allocator.

IMHO the main issue here is not extending live ranges too much. As for the runtime overhead, sinking a foldable address is assumed to not increase the execution time, because of the checks in isLegalAddressingMode and (with this patch) in canFoldAddrModeIntoCall.

I was thinking more in terms of sinking to arbitrary uses, as opposed to only sinking some uses.

[chill]

Ping?

[efriedma]

For constructs like the given testcase, you don't need to reason about whether it's profitable to fold into a call; cloning a GEP into both sides of an if-else doesn't actually increase execution time, so it's obviously profitable even if you can only fold it on one side of the if-else. Can you add some examples that actually require predicting whether the GEP is as cheap as a mov, and the argument is passed in a register?

llvm/include/llvm/CodeGen/TargetLowering.h

3178

In this case, the addressing mode doesn't actually represent any computation, so it isn't relevant for this transform; when do you expect it to become relevant?

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

24528

Counting arguments like this won't do what you want if there are arguments which are passed in multiple registers.

[chill]

llvm/include/llvm/CodeGen/TargetLowering.h

3178

The logic of the AddressingModeMatcher is that it accumulates an addressing mode from partial "expressions". For example an addressing mode
which ends up as [reg + imm] or [reg + reg] would have involved checks for legality of just [reg] here:
https://github.com/llvm/llvm-project/blob/ddfb974d0fca62e3eaeb98b79b5e29738c9082d2/llvm/lib/CodeGen/CodeGenPrepare.cpp#L4926

[chill]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

24528

Indeed.

For now, I'm proposing a different approach, doing the transformation entirely in MachineSink.
The relevant review is https://reviews.llvm.org/D152828

If it doesn't fly, I'll come back to this review and think how to solve this above issue.

[chill]

In D143898#4303461, @efriedma wrote:

For constructs like the given testcase, you don't need to reason about whether it's profitable to fold into a call; cloning a GEP into both sides of an if-else doesn't actually increase execution time, so it's obviously profitable even if you can only fold it on one side of the if-else. Can you add some examples that actually require predicting whether the GEP is as cheap as a mov, and the argument is passed in a register?

These test cases are constructed just to have the addressing computation and its users in separate basic blocks.
In our motivating examples, we have the addressing computation as loop invariant and its users inside a loop body, so
just sinking copies might increase run time.

I have added such test cases in https://reviews.llvm.org/D152828 : sink-and-fold.ll, functions f4 and f5.

llvm/include/llvm/CodeGen/TargetLowering.h

public:
   /// returned.
   virtual Value *createComplexDeinterleavingIR(
       Instruction *I, ComplexDeinterleavingOperation OperationType,
       ComplexDeinterleavingRotation Rotation, Value *InputA, Value *InputB,
       Value *Accumulator = nullptr) const {
     return nullptr;
   }
 
+  /// Return true if it is beneficial to fold address calculation into a call
+  /// sequence.
+  virtual bool canFoldAddrModeIntoCall(const CallBase &CB, unsigned ArgNo,
+                                       const AddrMode &AM) const {
+    return AM.HasBaseReg && AM.BaseGV == nullptr && AM.Scale == 0 && AM.BaseOffs == 0;
+  }
+
   //===--------------------------------------------------------------------===//
   // Runtime Library hooks
   //
 
   /// Rename the default libcall routine name for the specified libcall.
   void setLibcallName(RTLIB::Libcall Call, const char *Name) {
     LibcallRoutineNames[Call] = Name;
   }

llvm/include/llvm/CodeGen/TargetSubtargetInfo.h

public:
   virtual bool enablePostRAMachineScheduler() const;
 
   /// True if the subtarget should run the atomic expansion pass.
   virtual bool enableAtomicExpand() const;
 
   /// True if the subtarget should run the indirectbr expansion pass.
   virtual bool enableIndirectBrExpand() const;
 
+  /// Enable folding of address computations into call sequences.
+  virtual bool enableCallAddrFold() const { return false; }
+
   /// Override generic scheduling policy within a region.
   ///
   /// This is a convenient way for targets that don't provide any custom
   /// scheduling heuristics (no custom MachineSchedStrategy) to make
   /// changes to the generic scheduling policy.
   virtual void overrideSchedPolicy(MachineSchedPolicy &Policy,
                                    unsigned NumRegionInstrs) const {}
 

llvm/lib/CodeGen/CodeGenPrepare.cpp

static bool despeculateCountZeros(IntrinsicInst *CountZeros,
   // We are explicitly handling the zero case, so we can set the intrinsic's
   // undefined zero argument to 'true'. This will also prevent reprocessing the
   // intrinsic; we only despeculate when a zero input is defined.
   CountZeros->setArgOperand(1, Builder.getTrue());
   ModifiedDT = ModifyDT::ModifyBBDT;
   return true;
 }
 
+static bool isRegularCall(const CallBase *CB) {
+  return !CB->isInlineAsm() && CB->getIntrinsicID() == Intrinsic::not_intrinsic;
+}
+
+static const CallBase *getRegularCall(const Instruction *I) {
+  if (const auto *CB = dyn_cast<CallBase>(I); CB && isRegularCall(CB))
+    return CB;
+  return nullptr;
+}
+
+static unsigned getRegularCallArgNo(const Instruction *I, Value *V) {
+  if (const CallBase *CB = getRegularCall(I)) {
+    unsigned AN = 0;
+    for (const Use &U : CB->args()) {
+      if (U == V)
+        return AN;
+      ++AN;
+    }
+  }
+  return ~0;
+}
+
 bool CodeGenPrepare::optimizeCallInst(CallInst *CI, ModifyDT &ModifiedDT) {
   BasicBlock *BB = CI->getParent();
 
   // Lower inline assembly if we can.
   // If we found an inline asm expession, and if the target knows how to
   // lower it to normal LLVM code, do so now.
   if (CI->isInlineAsm()) {
     if (TLI->ExpandInlineAsm(CI)) {
if (MemTransferInst *MTI = dyn_cast<MemTransferInst>(MI)) {
         MTI->setSourceAlignment(SrcAlign);
     }
   }
 
   // If we have a cold call site, try to sink addressing computation into the
   // cold block.  This interacts with our handling for loads and stores to
   // ensure that we can fold all uses of a potential addressing computation
   // into their uses.  TODO: generalize this to work over profiling data
-  if (CI->hasFnAttr(Attribute::Cold) && !OptSize &&
-      !llvm::shouldOptimizeForSize(BB, PSI, BFI.get()))
+  if (isRegularCall(CI) && CI->hasFnAttr(Attribute::Cold) && !OptSize &&
+      !llvm::shouldOptimizeForSize(BB, PSI, BFI.get())) {
     for (auto &Arg : CI->args()) {
       if (!Arg->getType()->isPointerTy())
         continue;
       unsigned AS = Arg->getType()->getPointerAddressSpace();
       if (optimizeMemoryInst(CI, Arg, Arg->getType(), AS))
         return true;
     }
+  }
 
   IntrinsicInst *II = dyn_cast<IntrinsicInst>(CI);
   if (II) {
     switch (II->getIntrinsicID()) {
     default:
       break;
     case Intrinsic::assume:
       llvm_unreachable("llvm.assume should have been removed already");

 }
 
 namespace {
 
 /// A helper class for matching addressing modes.
 ///
 /// This encapsulates the logic for matching the target-legal addressing modes.
 class AddressingModeMatcher {
+  static constexpr unsigned INVALID_ARG_NO = ~0u;
+
   SmallVectorImpl<Instruction *> &AddrModeInsts;
   const TargetLowering &TLI;
+  const TargetSubtargetInfo &STI;
   const TargetRegisterInfo &TRI;
   const DataLayout &DL;
   const LoopInfo &LI;
   const std::function<const DominatorTree &()> getDTFn;
 
   /// AccessTy/MemoryInst - This is the type for the access (e.g. double) and
   /// the memory instruction that we're computing this address for.
   Type *AccessTy;
   unsigned AddrSpace;
+  unsigned ArgNo;
   Instruction *MemoryInst;
 
   /// This is the addressing mode that we're building up. This is
   /// part of the return value of this addressing mode matching stuff.
   ExtAddrMode &AddrMode;
 
   /// The instructions inserted by other CodeGenPrepare optimizations.
   const SetOfInstrs &InsertedInsts;
class AddressingModeMatcher {
   /// True if we are optimizing for size.
   bool OptSize;
 
   ProfileSummaryInfo *PSI;
   BlockFrequencyInfo *BFI;
 
   AddressingModeMatcher(
       SmallVectorImpl<Instruction *> &AMI, const TargetLowering &TLI,
+      const TargetSubtargetInfo &STI,
       const TargetRegisterInfo &TRI, const LoopInfo &LI,
       const std::function<const DominatorTree &()> getDTFn, Type *AT,
       unsigned AS, Instruction *MI, ExtAddrMode &AM,
       const SetOfInstrs &InsertedInsts, InstrToOrigTy &PromotedInsts,
       TypePromotionTransaction &TPT,
       std::pair<AssertingVH<GetElementPtrInst>, int64_t> &LargeOffsetGEP,
       bool OptSize, ProfileSummaryInfo *PSI, BlockFrequencyInfo *BFI)
-      : AddrModeInsts(AMI), TLI(TLI), TRI(TRI),
+      : AddrModeInsts(AMI), TLI(TLI), STI(STI), TRI(TRI),
         DL(MI->getModule()->getDataLayout()), LI(LI), getDTFn(getDTFn),
-        AccessTy(AT), AddrSpace(AS), MemoryInst(MI), AddrMode(AM),
-        InsertedInsts(InsertedInsts), PromotedInsts(PromotedInsts), TPT(TPT),
-        LargeOffsetGEP(LargeOffsetGEP), OptSize(OptSize), PSI(PSI), BFI(BFI) {
+        AccessTy(AT), AddrSpace(AS), ArgNo(INVALID_ARG_NO), MemoryInst(MI),
+        AddrMode(AM), InsertedInsts(InsertedInsts),
+        PromotedInsts(PromotedInsts), TPT(TPT), LargeOffsetGEP(LargeOffsetGEP),
+        OptSize(OptSize), PSI(PSI), BFI(BFI) {
     IgnoreProfitability = false;
   }
 
 public:
   /// Find the maximal addressing mode that a load/store of V can fold,
   /// give an access type of AccessTy.  This returns a list of involved
   /// instructions in AddrModeInsts.
   /// \p InsertedInsts The instructions inserted by other CodeGenPrepare
   /// optimizations.
   /// \p PromotedInsts maps the instructions to their type before promotion.
   /// \p The ongoing transaction where every action should be registered.
   static ExtAddrMode
   Match(Value *V, Type *AccessTy, unsigned AS, Instruction *MemoryInst,
         SmallVectorImpl<Instruction *> &AddrModeInsts,
         const TargetLowering &TLI, const LoopInfo &LI,
         const std::function<const DominatorTree &()> getDTFn,
+        const TargetSubtargetInfo &STI,
         const TargetRegisterInfo &TRI, const SetOfInstrs &InsertedInsts,
         InstrToOrigTy &PromotedInsts, TypePromotionTransaction &TPT,
         std::pair<AssertingVH<GetElementPtrInst>, int64_t> &LargeOffsetGEP,
         bool OptSize, ProfileSummaryInfo *PSI, BlockFrequencyInfo *BFI) {
     ExtAddrMode Result;
 
-    bool Success = AddressingModeMatcher(AddrModeInsts, TLI, TRI, LI, getDTFn,
-                                         AccessTy, AS, MemoryInst, Result,
-                                         InsertedInsts, PromotedInsts, TPT,
-                                         LargeOffsetGEP, OptSize, PSI, BFI)
-                       .matchAddr(V, 0);
+    AddressingModeMatcher Matcher(
+        AddrModeInsts, TLI, STI, TRI, LI, getDTFn, AccessTy, AS, MemoryInst, Result,
+        InsertedInsts, PromotedInsts, TPT, LargeOffsetGEP, OptSize, PSI, BFI);
+    if (STI.enableCallAddrFold())
+      Matcher.ArgNo = getRegularCallArgNo(MemoryInst, V);
+    bool Success = Matcher.matchAddr(V, 0);
     (void)Success;
     assert(Success && "Couldn't select *anything*?");
     return Result;
   }
 
 private:
   bool matchScaledValue(Value *ScaleReg, int64_t Scale, unsigned Depth);
   bool matchAddr(Value *Addr, unsigned Depth);
   bool matchOperationAddr(User *AddrInst, unsigned Opcode, unsigned Depth,
                           bool *MovedAway = nullptr);
   bool isProfitableToFoldIntoAddressingMode(Instruction *I,
                                             ExtAddrMode &AMBefore,
                                             ExtAddrMode &AMAfter);
   bool valueAlreadyLiveAtInst(Value *Val, Value *KnownLive1, Value *KnownLive2);
   bool isPromotionProfitable(unsigned NewCost, unsigned OldCost,
                              Value *PromotedOperand) const;
+  bool canFoldAddr(const ExtAddrMode &TestAddrMode) const;
 };
 
 class PhiNodeSet;
 
 /// An iterator for PhiNodeSet.
 class PhiNodeSetIterator {
   PhiNodeSet *const Set;
   size_t CurrentIndex = 0;
bool AddressingModeMatcher::matchScaledValue(Value *ScaleReg, int64_t Scale,
   ExtAddrMode TestAddrMode = AddrMode;
 
   // Add scale to turn X*4+X*3 -> X*7.  This could also do things like
   // [A+B + A*7] -> [B+A*8].
   TestAddrMode.Scale += Scale;
   TestAddrMode.ScaledReg = ScaleReg;
 
   // If the new address isn't legal, bail out.
-  if (!TLI.isLegalAddressingMode(DL, TestAddrMode, AccessTy, AddrSpace))
+  if (!canFoldAddr(TestAddrMode))
     return false;
 
   // It was legal, so commit it.
   AddrMode = TestAddrMode;
 
   // Okay, we decided that we can add ScaleReg+Scale to AddrMode.  Check now
   // to see if ScaleReg is actually X+C.  If so, we can turn this into adding
   // X*Scale + C*Scale to addr mode. If we found available IV increment, do not
   // go any further: we can reuse it and cannot eliminate it.
   ConstantInt *CI = nullptr;
   Value *AddLHS = nullptr;
   if (isa<Instruction>(ScaleReg) && // not a constant expr.
       match(ScaleReg, m_Add(m_Value(AddLHS), m_ConstantInt(CI))) &&
       !isIVIncrement(ScaleReg, &LI) && CI->getValue().isSignedIntN(64)) {
     TestAddrMode.InBounds = false;
     TestAddrMode.ScaledReg = AddLHS;
     TestAddrMode.BaseOffs += CI->getSExtValue() * TestAddrMode.Scale;
 
     // If this addressing mode is legal, commit it and remember that we folded
     // this instruction.
-    if (TLI.isLegalAddressingMode(DL, TestAddrMode, AccessTy, AddrSpace)) {
+    if (canFoldAddr(TestAddrMode)) {
       AddrModeInsts.push_back(cast<Instruction>(ScaleReg));
       AddrMode = TestAddrMode;
       return true;
     }
     // Restore status quo.
     TestAddrMode = AddrMode;
   }
 
if (auto IVStep = GetConstantStep(ScaleReg)) {
       APInt Offset = Step * AddrMode.Scale;
       if (Offset.isSignedIntN(64)) {
         TestAddrMode.InBounds = false;
         TestAddrMode.ScaledReg = IVInc;
         TestAddrMode.BaseOffs -= Offset.getLimitedValue();
         // If this addressing mode is legal, commit it..
         // (Note that we defer the (expensive) domtree base legality check
         // to the very last possible point.)
-        if (TLI.isLegalAddressingMode(DL, TestAddrMode, AccessTy, AddrSpace) &&
+        if (canFoldAddr(TestAddrMode) &&
             getDTFn().dominates(IVInc, MemoryInst)) {
           AddrModeInsts.push_back(cast<Instruction>(IVInc));
           AddrMode = TestAddrMode;
           return true;
         }
         // Restore status quo.
         TestAddrMode = AddrMode;
       }
bool AddressingModeMatcher::isPromotionProfitable(
   if (NewCost < OldCost)
     return true;
   // The promotion is neutral but it may help folding the sign extension in
   // loads for instance.
   // Check that we did not create an illegal instruction.
   return isPromotedInstructionLegal(TLI, DL, PromotedOperand);
 }
 
+bool AddressingModeMatcher::canFoldAddr(const ExtAddrMode &TestAddrMode) const {
+  if (ArgNo != INVALID_ARG_NO) {
+    // Check if the address is "foldable" into a regular call.
+    const auto &CB = cast<CallBase>(*MemoryInst);
+    if (STI.enableCallAddrFold() &&
+        TLI.canFoldAddrModeIntoCall(CB, ArgNo, TestAddrMode))
+      return true;
+
+    // Even if it isn't, accept it if we have a cold call, but still require
+    // legal addressing mode. This limits the amount of code we potentially
+    // sink.
+    if (!CB.hasFnAttr(Attribute::Cold) || OptSize ||
+        llvm::shouldOptimizeForSize(CB.getParent(), PSI, BFI))
+      return false;
+  }
+
+  return TLI.isLegalAddressingMode(DL, TestAddrMode, AccessTy, AddrSpace);
+}
+
 /// Given an instruction or constant expr, see if we can fold the operation
 /// into the addressing mode. If so, update the addressing mode and return
 /// true, otherwise return false without modifying AddrMode.
 /// If \p MovedAway is not NULL, it contains the information of whether or
 /// not AddrInst has to be folded into the addressing mode on success.
 /// If \p MovedAway == true, \p AddrInst will not be part of the addressing
 /// because it has been moved away.
 /// Thus AddrInst must not be added in the matched instructions.
for (unsigned i = 1, e = AddrInst->getNumOperands(); i != e; ++i, ++GTI) {
           VariableScale = TypeSize;
         }
       }
     }
 
     // A common case is for the GEP to only do a constant offset.  In this case,
     // just add it to the disp field and check validity.
     if (VariableOperand == -1) {
+      bool InBounds = AddrMode.InBounds;
+      int64_t BaseOffs = AddrMode.BaseOffs;
       AddrMode.BaseOffs += ConstantOffset;
-      if (matchAddr(AddrInst->getOperand(0), Depth + 1)) {
           if (!cast<GEPOperator>(AddrInst)->isInBounds())
             AddrMode.InBounds = false;
+      if (matchAddr(AddrInst->getOperand(0), Depth + 1))
           return true;
-      }
-      AddrMode.BaseOffs -= ConstantOffset;
+      AddrMode.InBounds = InBounds;
+      AddrMode.BaseOffs = BaseOffs;
 
       if (EnableGEPOffsetSplit && isa<GetElementPtrInst>(AddrInst) &&
           TLI.shouldConsiderGEPOffsetSplit() && Depth == 0 &&
           ConstantOffset > 0) {
           // Record GEPs with non-zero offsets as candidates for splitting in
           // the event that the offset cannot fit into the r+i addressing mode.
           // Simple and common case that only one GEP is used in calculating the
           // address for the memory access.
bool AddressingModeMatcher::matchAddr(Value *Addr, unsigned Depth) {
   // Start a transaction at this point that we will rollback if the matching
   // fails.
   TypePromotionTransaction::ConstRestorationPt LastKnownGood =
       TPT.getRestorationPoint();
   if (ConstantInt *CI = dyn_cast<ConstantInt>(Addr)) {
     if (CI->getValue().isSignedIntN(64)) {
       // Fold in immediates if legal for the target.
       AddrMode.BaseOffs += CI->getSExtValue();
-      if (TLI.isLegalAddressingMode(DL, AddrMode, AccessTy, AddrSpace))
+      if (canFoldAddr(AddrMode))
         return true;
       AddrMode.BaseOffs -= CI->getSExtValue();
     }
   } else if (GlobalValue *GV = dyn_cast<GlobalValue>(Addr)) {
     // If this is a global variable, try to fold it into the addressing mode.
     if (!AddrMode.BaseGV) {
       AddrMode.BaseGV = GV;
-      if (TLI.isLegalAddressingMode(DL, AddrMode, AccessTy, AddrSpace))
+      if (canFoldAddr(AddrMode))
         return true;
       AddrMode.BaseGV = nullptr;
     }
   } else if (Instruction *I = dyn_cast<Instruction>(Addr)) {
     ExtAddrMode BackupAddrMode = AddrMode;
     unsigned OldSize = AddrModeInsts.size();
 
     // Check to see if it is possible to fold this operation.
if (ConstantInt *CI = dyn_cast<ConstantInt>(Addr)) {
     return true;
   }
 
   // Worse case, the target should support [reg] addressing modes. :)
   if (!AddrMode.HasBaseReg) {
     AddrMode.HasBaseReg = true;
     AddrMode.BaseReg = Addr;
     // Still check for legality in case the target supports [imm] but not [i+r].
-    if (TLI.isLegalAddressingMode(DL, AddrMode, AccessTy, AddrSpace))
+    if (canFoldAddr(AddrMode))
       return true;
     AddrMode.HasBaseReg = false;
     AddrMode.BaseReg = nullptr;
   }
 
   // If the base register is already taken, see if we can do [r+r].
   if (AddrMode.Scale == 0) {
     AddrMode.Scale = 1;
     AddrMode.ScaledReg = Addr;
-    if (TLI.isLegalAddressingMode(DL, AddrMode, AccessTy, AddrSpace))
+    if (canFoldAddr(AddrMode))
       return true;
     AddrMode.Scale = 0;
     AddrMode.ScaledReg = nullptr;
   }
   // Couldn't match.
   TPT.rollback(LastKnownGood);
   return false;
 }

 static constexpr int MaxMemoryUsesToScan = 100;
 
 /// Recursively walk all the uses of I until we find a memory use.
 /// If we find an obviously non-foldable instruction, return true.
 /// Add accessed addresses and types to MemoryUses.
 static bool FindAllMemoryUses(
     Instruction *I, SmallVectorImpl<std::pair<Use *, Type *>> &MemoryUses,
     SmallPtrSetImpl<Instruction *> &ConsideredInsts, const TargetLowering &TLI,
+    const TargetSubtargetInfo &STI,
     const TargetRegisterInfo &TRI, bool OptSize, ProfileSummaryInfo *PSI,
     BlockFrequencyInfo *BFI, int &SeenInsts) {
   // If we already considered this instruction, we're done.
   if (!ConsideredInsts.insert(I).second)
     return false;
 
   // If this is an obviously unfoldable instruction, bail out.
   if (!MightBeFoldableInst(I))
if (CallInst *CI = dyn_cast<CallInst>(UserI)) {
         // If this is a cold call, we can sink the addressing calculation into
         // the cold path.  See optimizeCallInst
         bool OptForSize =
             OptSize || llvm::shouldOptimizeForSize(CI->getParent(), PSI, BFI);
         if (!OptForSize)
           continue;
       }
 
-      InlineAsm *IA = dyn_cast<InlineAsm>(CI->getCalledOperand());
-      if (!IA)
-        return true;
-
+      if (InlineAsm *IA = dyn_cast<InlineAsm>(CI->getCalledOperand())) {
         // If this is a memory operand, we're cool, otherwise bail out.
         if (!IsOperandAMemoryOperand(CI, IA, I, TLI, TRI))
           return true;
         continue;
       }
 
-    if (FindAllMemoryUses(UserI, MemoryUses, ConsideredInsts, TLI, TRI, OptSize,
+      // Bail if call folding is not enabled.
+      if (!STI.enableCallAddrFold())
+        return true;
+
+      // Intrinsics are handled elsewhere and we can't quite handle non-pointer
+      // types yet.
+      if (isa<IntrinsicInst>(UserI) || !I->getType()->isPointerTy())
+        return true;
+
+      MemoryUses.push_back({&U, I->getType()});
+      continue;
+    }
+
+    if (FindAllMemoryUses(UserI, MemoryUses, ConsideredInsts, TLI, STI, TRI, OptSize,
                           PSI, BFI, SeenInsts))
       return true;
   }
 
   return false;
 }
 
 static bool FindAllMemoryUses(
     Instruction *I, SmallVectorImpl<std::pair<Use *, Type *>> &MemoryUses,
-    const TargetLowering &TLI, const TargetRegisterInfo &TRI, bool OptSize,
+    const TargetLowering &TLI, const TargetSubtargetInfo &STI, const TargetRegisterInfo &TRI, bool OptSize,
     ProfileSummaryInfo *PSI, BlockFrequencyInfo *BFI) {
   int SeenInsts = 0;
   SmallPtrSet<Instruction *, 16> ConsideredInsts;
-  return FindAllMemoryUses(I, MemoryUses, ConsideredInsts, TLI, TRI, OptSize,
+  return FindAllMemoryUses(I, MemoryUses, ConsideredInsts, TLI, STI, TRI, OptSize,
                            PSI, BFI, SeenInsts);
 }
 
 static unsigned MaxLoopInvariantUsesToScan = 20;
 static bool isUsedInLoop(const Value *V, const Loop *L) {
   unsigned N = 0;
 
   for (const Use &U : V->uses()) {
bool AddressingModeMatcher::isProfitableToFoldIntoAddressingMode(
   if (!BaseReg && !ScaledReg)
     return true;
 
   // If all uses of this instruction can have the address mode sunk into them,
   // we can remove the addressing mode and effectively trade one live register
   // for another (at worst.)  In this context, folding an addressing mode into
   // the use is just a particularly nice way of sinking it.
   SmallVector<std::pair<Use *, Type *>, 16> MemoryUses;
-  if (FindAllMemoryUses(I, MemoryUses, TLI, TRI, OptSize, PSI, BFI))
+  if (FindAllMemoryUses(I, MemoryUses, TLI, STI, TRI, OptSize, PSI, BFI))
     return false; // Has a non-memory, non-foldable use!
 
   // Now that we know that all uses of this instruction are part of a chain of
   // computation involving only operations that could theoretically be folded
   // into a memory use, loop over each of these memory operation uses and see
   // if they could  *actually* fold the instruction.  The assumption is that
   // addressing modes are cheap and that duplicating the computation involved
   // many times is worthwhile, even on a fastpath. For sinking candidates
for (const std::pair<Use *, Type *> &Pair : MemoryUses) {
     // Do a match against the root of this address, ignoring profitability. This
     // will tell us if the addressing mode for the memory operation will
     // *actually* cover the shared instruction.
     ExtAddrMode Result;
     std::pair<AssertingVH<GetElementPtrInst>, int64_t> LargeOffsetGEP(nullptr,
                                                                       0);
     TypePromotionTransaction::ConstRestorationPt LastKnownGood =
         TPT.getRestorationPoint();
-    AddressingModeMatcher Matcher(MatchedAddrModeInsts, TLI, TRI, LI, getDTFn,
+    AddressingModeMatcher Matcher(MatchedAddrModeInsts, TLI, STI, TRI, LI, getDTFn,
                                   AddressAccessTy, AS, UserI, Result,
                                   InsertedInsts, PromotedInsts, TPT,
                                   LargeOffsetGEP, OptSize, PSI, BFI);
     Matcher.IgnoreProfitability = true;
+    if (STI.enableCallAddrFold())
+      Matcher.ArgNo = getRegularCallArgNo(UserI, Address);
     bool Success = Matcher.matchAddr(Address, 0);
     (void)Success;
     assert(Success && "Couldn't select *anything*?");
 
     // The match was to check the profitability, the changes made are not
     // part of the original matcher. Therefore, they should be dropped
     // otherwise the original matcher will not present the right state.
     TPT.rollback(LastKnownGood);
while (!worklist.empty()) {
     // actual use.  It's expected that most address matches don't actually need
     // the domtree.
     auto getDTFn = [MemoryInst, this]() -> const DominatorTree & {
       Function *F = MemoryInst->getParent()->getParent();
       return this->getDT(*F);
     };
     ExtAddrMode NewAddrMode = AddressingModeMatcher::Match(
         V, AccessTy, AddrSpace, MemoryInst, AddrModeInsts, *TLI, *LI, getDTFn,
-        *TRI, InsertedInsts, PromotedInsts, TPT, LargeOffsetGEP, OptSize, PSI,
+        *SubtargetInfo, *TRI, InsertedInsts, PromotedInsts, TPT, LargeOffsetGEP, OptSize, PSI,
         BFI.get());
 
     GetElementPtrInst *GEP = LargeOffsetGEP.first;
     if (GEP && !NewGEPBases.count(GEP)) {
       // If splitting the underlying data structure can reduce the offset of a
       // GEP, collect the GEP.  Skip the GEPs that are the new bases of
       // previously split data structures.
       LargeOffsetGEPMap[GEP->getPointerOperand()].push_back(LargeOffsetGEP);

llvm/lib/Target/AArch64/AArch64ISelLowering.h

public:
   bool isComplexDeinterleavingOperationSupported(
       ComplexDeinterleavingOperation Operation, Type *Ty) const override;
 
   Value *createComplexDeinterleavingIR(
       Instruction *I, ComplexDeinterleavingOperation OperationType,
       ComplexDeinterleavingRotation Rotation, Value *InputA, Value *InputB,
       Value *Accumulator = nullptr) const override;
 
+  bool canFoldAddrModeIntoCall(const CallBase &CB, unsigned ArgNo,
+                               const AddrMode &AM) const override;
+
   bool supportSplitCSR(MachineFunction *MF) const override {
     return MF->getFunction().getCallingConv() == CallingConv::CXX_FAST_TLS &&
            MF->getFunction().hasFnAttribute(Attribute::NoUnwind);
   }
   void initializeSplitCSR(MachineBasicBlock *Entry) const override;
   void insertCopiesSplitCSR(
       MachineBasicBlock *Entry,
       const SmallVectorImpl<MachineBasicBlock *> &Exits) const override;

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (OperationType == ComplexDeinterleavingOperation::CAdd) {
     if (IntId == Intrinsic::not_intrinsic)
       return nullptr;
 
     return B.CreateIntrinsic(IntId, Ty, {InputA, InputB});
   }
 
   return nullptr;
 }
+
+bool AArch64TargetLowering::canFoldAddrModeIntoCall(const CallBase &CB,
+                                                    unsigned ArgNo,
+                                                    const AddrMode &AM) const {
+  // We should always accept a single base register.
+  if (TargetLowering::canFoldAddrModeIntoCall(CB, ArgNo, AM))
+    return true;
+
+  if (ArgNo > 7 || AM.BaseGV || !AM.HasBaseReg)
+    return false;
+
+  // For more complex addressing modes, check the possibility of a cheap
+  // materialisation into an argument register.
+
+  // reg + imm
+  if (AM.Scale == 0)
+    return isLegalAddImmediate(AM.BaseOffs);
+
+  if (AM.BaseOffs != 0)
+    return false;
+
+  // reg + scale * reg
+  if (AM.Scale == 1)
+    return true;
+
+  // Some CPUs have fast `reg + scale * reg` instruction, for scales of 2, 4, 8, and 16.
+  if (!Subtarget->hasLSLFast() || AM.Scale <= 0)
+    return false;
+
+  uint64_t S = uint64_t(AM.Scale);
+  return (S & (S - 1)) == 0 && S <= 16;
+}

llvm/lib/Target/AArch64/AArch64Subtarget.h

#include "AArch64GenSubtargetInfo.inc"
   const CallLowering *getCallLowering() const override;
   const InlineAsmLowering *getInlineAsmLowering() const override;
   InstructionSelector *getInstructionSelector() const override;
   const LegalizerInfo *getLegalizerInfo() const override;
   const RegisterBankInfo *getRegBankInfo() const override;
   const Triple &getTargetTriple() const { return TargetTriple; }
   bool enableMachineScheduler() const override { return true; }
   bool enablePostRAScheduler() const override { return usePostRAScheduler(); }
+  bool enableCallAddrFold() const override { return true; }
   
   /// Returns ARM processor family.
   /// Avoid this function! CPU specifics should be kept local to this class
   /// and preferably modeled with SubtargetFeatures or properties in
   /// initializeProperties().
   ARMProcFamilyEnum getProcFamily() const {
     return ARMProcFamily;
   }
 

llvm/test/CodeGen/AArch64/call-addr-fold.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -codegenprepare < %s | FileCheck %s
+target triple = "aarch64-linux"
+
+declare void @use(...)
+
+; Check address is sunk towards the load, since CodeGenPrepare considers it likely to be
+; "folded" into the call as well.
+define i32 @f0(i1 %c1, ptr %p) {
+; CHECK-LABEL: @f0(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[A:%.*]] = getelementptr i32, ptr [[P:%.*]], i32 2
+; CHECK-NEXT:    br i1 [[C1:%.*]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    [[V0:%.*]] = call i32 @use(ptr [[A]])
+; CHECK-NEXT:    br label [[EXIT:%.*]]
+; CHECK:       if.else:
+; CHECK-NEXT:    [[SUNKADDR:%.*]] = getelementptr i8, ptr [[P]], i64 8
+; CHECK-NEXT:    [[V1:%.*]] = load i32, ptr [[SUNKADDR]], align 4
+; CHECK-NEXT:    br label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[V:%.*]] = phi i32 [ [[V0]], [[IF_THEN]] ], [ [[V1]], [[IF_ELSE]] ]
+; CHECK-NEXT:    ret i32 [[V]]
+;
+entry:
+  %a = getelementptr i32, ptr %p, i32 2
+  br i1 %c1, label %if.then, label %if.else
+
+if.then:
+  %v0 = call i32 @use(ptr %a)
+  br label %exit
+
+if.else:
+  %v1 = load i32, ptr %a
+  br label %exit
+
+exit:
+  %v = phi i32 [%v0, %if.then], [%v1, %if.else]
+  ret i32 %v
+}
+
+define i32 @f1(i1 %c1, ptr %p, i64 %i) {
+; CHECK-LABEL: @f1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[A:%.*]] = getelementptr i8, ptr [[P:%.*]], i64 [[I:%.*]]
+; CHECK-NEXT:    br i1 [[C1:%.*]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    [[V0:%.*]] = call i32 @use(ptr [[A]])
+; CHECK-NEXT:    br label [[EXIT:%.*]]
+; CHECK:       if.else:
+; CHECK-NEXT:    [[SUNKADDR:%.*]] = getelementptr i8, ptr [[P]], i64 [[I]]
+; CHECK-NEXT:    [[V1:%.*]] = load i32, ptr [[SUNKADDR]], align 4
+; CHECK-NEXT:    br label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[V:%.*]] = phi i32 [ [[V0]], [[IF_THEN]] ], [ [[V1]], [[IF_ELSE]] ]
+; CHECK-NEXT:    ret i32 [[V]]
+;
+entry:
+  %a = getelementptr i8, ptr %p, i64 %i
+  br i1 %c1, label %if.then, label %if.else
+
+if.then:
+  %v0 = call i32 @use(ptr %a)
+  br label %exit
+
+if.else:
+  %v1 = load i32, ptr %a
+  br label %exit
+
+exit:
+  %v = phi i32 [%v0, %if.then], [%v1, %if.else]
+  ret i32 %v
+}
+
+; Address calculation too complex.
+define i32 @f2(i1 %c1, ptr %p, i64 %i)  {
+; CHECK-LABEL: @f2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[A:%.*]] = getelementptr i32, ptr [[P:%.*]], i64 [[I:%.*]]
+; CHECK-NEXT:    br i1 [[C1:%.*]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    [[V0:%.*]] = call i32 @use(ptr [[A]])
+; CHECK-NEXT:    br label [[EXIT:%.*]]
+; CHECK:       if.else:
+; CHECK-NEXT:    [[V1:%.*]] = load i32, ptr [[A]], align 4
+; CHECK-NEXT:    br label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[V:%.*]] = phi i32 [ [[V0]], [[IF_THEN]] ], [ [[V1]], [[IF_ELSE]] ]
+; CHECK-NEXT:    ret i32 [[V]]
+;
+entry:
+  %a = getelementptr i32, ptr %p, i64 %i
+  br i1 %c1, label %if.then, label %if.else
+
+if.then:
+  %v0 = call i32 @use(ptr %a)
+  br label %exit
+
+if.else:
+  %v1 = load i32, ptr %a
+  br label %exit
+
+exit:
+  %v = phi i32 [%v0, %if.then], [%v1, %if.else]
+  ret i32 %v
+}
+
+; Address calculation cheap enough on some cores.
+define i32 @f3(i1 %c1, ptr %p, i64 %i) "target-cpu"="neoverse-n1" {
+; CHECK-LABEL: @f3(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[A:%.*]] = getelementptr i32, ptr [[P:%.*]], i64 [[I:%.*]]
+; CHECK-NEXT:    br i1 [[C1:%.*]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    [[V0:%.*]] = call i32 @use(ptr [[A]])
+; CHECK-NEXT:    br label [[EXIT:%.*]]
+; CHECK:       if.else:
+; CHECK-NEXT:    [[SUNKADDR:%.*]] = mul i64 [[I]], 4
+; CHECK-NEXT:    [[SUNKADDR1:%.*]] = getelementptr i8, ptr [[P]], i64 [[SUNKADDR]]
+; CHECK-NEXT:    [[V1:%.*]] = load i32, ptr [[SUNKADDR1]], align 4
+; CHECK-NEXT:    br label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[V:%.*]] = phi i32 [ [[V0]], [[IF_THEN]] ], [ [[V1]], [[IF_ELSE]] ]
+; CHECK-NEXT:    ret i32 [[V]]
+;
+entry:
+  %a = getelementptr i32, ptr %p, i64 %i
+  br i1 %c1, label %if.then, label %if.else
+
+if.then:
+  %v0 = call i32 @use(ptr %a)
+  br label %exit
+
+if.else:
+  %v1 = load i32, ptr %a
+  br label %exit
+
+exit:
+  %v = phi i32 [%v0, %if.then], [%v1, %if.else]
+  ret i32 %v
+}
+
+define i32 @f4(i1 %c1, ptr %p, i64 %i) {
+; CHECK-LABEL: @f4(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[A:%.*]] = getelementptr i8, ptr [[P:%.*]], i64 [[I:%.*]]
+; CHECK-NEXT:    br i1 [[C1:%.*]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    [[V0:%.*]] = call i32 @use(ptr [[A]])
+; CHECK-NEXT:    br label [[EXIT:%.*]]
+; CHECK:       if.else:
+; CHECK-NEXT:    [[V1:%.*]] = call i32 @use(i32 1, ptr [[A]])
+; CHECK-NEXT:    br label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[V:%.*]] = phi i32 [ [[V0]], [[IF_THEN]] ], [ [[V1]], [[IF_ELSE]] ]
+; CHECK-NEXT:    ret i32 [[V]]
+;
+entry:
+  %a = getelementptr i8, ptr %p, i64 %i
+  br i1 %c1, label %if.then, label %if.else
+
+if.then:
+  %v0 = call i32 @use(ptr %a)
+  br label %exit
+
+if.else:
+  %v1 = call i32 @use(i32 1, ptr %a)
+  br label %exit
+
+exit:
+  %v = phi i32 [%v0, %if.then], [%v1, %if.else]
+  ret i32 %v
+}