[AArch64] Fix incorrect `isLegalAddressingMode`

[chill]

Summary

AArch64TargetLowering::isLegalAddressingMode has a number of defects, including
accepting an addressing mode, which consists of only an immediate operand, or not checking
the offset range for an addressing mode in the form 1*ScaledReg + Offs.

This patch fixes the above issues.

[chill]

Fixed a few places in LSR where we called isLegalAddressingMode without a base register, even
though we meant an addressing mode with a base register.

[chill]

+ canonicalise 2*Reg into Reg + Reg

[chill]

Ping?

[dmgreen]

I can certainly imagine that we haven't got Addressing modes without a BaseReg entirely correct in the past and just assumed it was present. The change in IsSimplerBaseSCEVForTarget I added recently and look more correct now.

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

3185

Can you give more details why this is changed to HasBaseReg? It seems that isAlwaysFoldable already sets ScaledReg to 1 or -1, so this now tests for BaseReg + 1/-1*ScaledReg + IncOffset?

[chill]

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

3185

Right, it's incorrect.
However, I'd rather see if I can fix it by changing isAlwaysFoldable, that invents a scale existence without any external input.

To be continued ...

[dmgreen]

Thanks. LGTM

[chill]

Work around TargetLowering::AddrMode no longer an aggregate in C++20.

[dmgreen]

LGTM still.

llvm/lib/CodeGen/CodeGenPrepare.cpp

case Instruction::AddrSpaceCast: {
     unsigned SrcAS =
         AddrInst->getOperand(0)->getType()->getPointerAddressSpace();
     unsigned DestAS = AddrInst->getType()->getPointerAddressSpace();
     if (TLI.getTargetMachine().isNoopAddrSpaceCast(SrcAS, DestAS))
       return matchAddr(AddrInst->getOperand(0), Depth);
     return false;
   }
   case Instruction::Add: {
-    // Check to see if we can merge in the RHS then the LHS.  If so, we win.
+    // Check to see if we can merge in one operand, then the other.  If so, we
+    // win.
     ExtAddrMode BackupAddrMode = AddrMode;
     unsigned OldSize = AddrModeInsts.size();
     // Start a transaction at this point.
     // The LHS may match but not the RHS.
     // Therefore, we need a higher level restoration point to undo partially
     // matched operation.
     TypePromotionTransaction::ConstRestorationPt LastKnownGood =
         TPT.getRestorationPoint();
 
+    // Try to match an integer constant second to increase its chance of ending
+    // up in `BaseOffs`, resp. decrease its chance of ending up in `BaseReg`.
+    int First = 0, Second = 1;
+    if (isa<ConstantInt>(AddrInst->getOperand(First))
+      && !isa<ConstantInt>(AddrInst->getOperand(Second)))
+        std::swap(First, Second);
     AddrMode.InBounds = false;
-    if (matchAddr(AddrInst->getOperand(1), Depth + 1) &&
-        matchAddr(AddrInst->getOperand(0), Depth + 1))
+    if (matchAddr(AddrInst->getOperand(First), Depth + 1) &&
+        matchAddr(AddrInst->getOperand(Second), Depth + 1))
       return true;
 
     // Restore the old addr mode info.
     AddrMode = BackupAddrMode;
     AddrModeInsts.resize(OldSize);
     TPT.rollback(LastKnownGood);
 
-    // Otherwise this was over-aggressive.  Try merging in the LHS then the RHS.
-    if (matchAddr(AddrInst->getOperand(0), Depth + 1) &&
-        matchAddr(AddrInst->getOperand(1), Depth + 1))
+    // Otherwise this was over-aggressive.  Try merging operands in the opposite
+    // order.
+    if (matchAddr(AddrInst->getOperand(Second), Depth + 1) &&
+        matchAddr(AddrInst->getOperand(First), Depth + 1))
       return true;
 
     // Otherwise we definitely can't merge the ADD in.
     AddrMode = BackupAddrMode;
     AddrModeInsts.resize(OldSize);
     TPT.rollback(LastKnownGood);
     break;
   }
for (unsigned i = 1, e = AddrInst->getNumOperands(); i != e; ++i, ++GTI) {
         }
       }
     }
 
     // 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) {
       AddrMode.BaseOffs += ConstantOffset;
-      if (ConstantOffset == 0 ||
-          TLI.isLegalAddressingMode(DL, AddrMode, AccessTy, AddrSpace)) {
-        // Check to see if we can fold the base pointer in too.
-        if (matchAddr(AddrInst->getOperand(0), Depth + 1)) {
           if (!cast<GEPOperator>(AddrInst)->isInBounds())
             AddrMode.InBounds = false;
           return true;
-        }
-      } else 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.
+      AddrMode.BaseOffs -= ConstantOffset;
+
+      if (EnableGEPOffsetSplit && isa<GetElementPtrInst>(AddrInst) &&
+          // 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.
           Value *Base = AddrInst->getOperand(0);
           auto *BaseI = dyn_cast<Instruction>(Base);
           auto *GEP = cast<GetElementPtrInst>(AddrInst);
           if (isa<Argument>(Base) || isa<GlobalValue>(Base) ||
               (BaseI && !isa<CastInst>(BaseI) &&
                !isa<GetElementPtrInst>(BaseI))) {
             // Make sure the parent block allows inserting non-PHI instructions
             // before the terminator.
-          BasicBlock *Parent =
-              BaseI ? BaseI->getParent() : &GEP->getFunction()->getEntryBlock();
+            BasicBlock *Parent = BaseI ? BaseI->getParent()
+                                       : &GEP->getFunction()->getEntryBlock();
             if (!Parent->getTerminator()->isEHPad())
             LargeOffsetGEP = std::make_pair(GEP, ConstantOffset);
           }
       }
-      AddrMode.BaseOffs -= ConstantOffset;
+
       return false;
     }
 
     // Save the valid addressing mode in case we can't match.
     ExtAddrMode BackupAddrMode = AddrMode;
     unsigned OldSize = AddrModeInsts.size();
 
     // See if the scale and offset amount is valid for this target.
for (auto &Entry : LargeOffsetGEPMap) {
     Value *NewBaseGEP = nullptr;
 
     auto *LargeOffsetGEP = LargeOffsetGEPs.begin();
     while (LargeOffsetGEP != LargeOffsetGEPs.end()) {
       GetElementPtrInst *GEP = LargeOffsetGEP->first;
       int64_t Offset = LargeOffsetGEP->second;
       if (Offset != BaseOffset) {
         TargetLowering::AddrMode AddrMode;
+        AddrMode.HasBaseReg = true;
         AddrMode.BaseOffs = Offset - BaseOffset;
         // The result type of the GEP might not be the type of the memory
         // access.
         if (!TLI->isLegalAddressingMode(*DL, AddrMode,
                                         GEP->getResultElementType(),
                                         GEP->getAddressSpace())) {
           // We need to create a new base if the offset to the current base is
           // too large to fit into the addressing mode. So, a very large struct

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

 // immediates is the same as for an add or a sub.
 bool AArch64TargetLowering::isLegalICmpImmediate(int64_t Immed) const {
   return isLegalAddImmediate(Immed);
 }
 
 /// isLegalAddressingMode - Return true if the addressing mode represented
 /// by AM is legal for this target, for a load/store of the specified type.
 bool AArch64TargetLowering::isLegalAddressingMode(const DataLayout &DL,
-                                                  const AddrMode &AM, Type *Ty,
+                                                  const AddrMode &AMode, Type *Ty,
                                                   unsigned AS, Instruction *I) const {
   // AArch64 has five basic addressing modes:
   //  reg
   //  reg + 9-bit signed offset
   //  reg + SIZE_IN_BYTES * 12-bit unsigned offset
   //  reg1 + reg2
   //  reg + SIZE_IN_BYTES * reg
 
   // No global is ever allowed as a base.
-  if (AM.BaseGV)
+  if (AMode.BaseGV)
     return false;
 
   // No reg+reg+imm addressing.
-  if (AM.HasBaseReg && AM.BaseOffs && AM.Scale)
+  if (AMode.HasBaseReg && AMode.BaseOffs && AMode.Scale)
+    return false;
+
+  // Canonicalise `1*ScaledReg + imm` into `BaseReg + imm` and
+  // `2*ScaledReg` into `BaseReg + ScaledReg`
+  AddrMode AM = AMode;
+  if (AM.Scale && !AM.HasBaseReg) {
+    if (AM.Scale == 1) {
+      AM.HasBaseReg = true;
+      AM.Scale = 0;
+    } else if (AM.Scale == 2) {
+      AM.HasBaseReg = true;
+      AM.Scale = 1;
+    } else {
+      return false;
+    }
+  }
+
+  // A base register is required in all addressing modes.
+  if (!AM.HasBaseReg)
     return false;
 
   // FIXME: Update this method to support scalable addressing modes.
   if (Ty->isScalableTargetExtTy())
     return AM.HasBaseReg && !AM.BaseOffs && !AM.Scale;
 
   if (isa<ScalableVectorType>(Ty)) {
     uint64_t VecElemNumBytes =

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

if (LatchBlock != ExitingBlock)
               goto decline_post_inc;
             // Check for possible scaled-address reuse.
             if (isAddressUse(TTI, UI->getUser(), UI->getOperandValToReplace())) {
               MemAccessTy AccessTy = getAccessType(
                   TTI, UI->getUser(), UI->getOperandValToReplace());
               int64_t Scale = C->getSExtValue();
               if (TTI.isLegalAddressingMode(AccessTy.MemTy, /*BaseGV=*/nullptr,
                                             /*BaseOffset=*/0,
-                                            /*HasBaseReg=*/false, Scale,
+                                            /*HasBaseReg=*/true, Scale,
                                             AccessTy.AddrSpace))
                 goto decline_post_inc;
               Scale = -Scale;
               if (TTI.isLegalAddressingMode(AccessTy.MemTy, /*BaseGV=*/nullptr,
                                             /*BaseOffset=*/0,
-                                            /*HasBaseReg=*/false, Scale,
+                                            /*HasBaseReg=*/true, Scale,
                                             AccessTy.AddrSpace))
                 goto decline_post_inc;
             }
           }
         }
 
     LLVM_DEBUG(dbgs() << "  Change loop exiting icmp to use postinc iv: "
                       << *Cond << '\n');
if (!IncConst || !isAddressUse(TTI, UserInst, Operand))
     return false;
 
   if (IncConst->getAPInt().getSignificantBits() > 64)
     return false;
 
   MemAccessTy AccessTy = getAccessType(TTI, UserInst, Operand);
   int64_t IncOffset = IncConst->getValue()->getSExtValue();
   if (!isAlwaysFoldable(TTI, LSRUse::Address, AccessTy, /*BaseGV=*/nullptr,
                         IncOffset, /*HasBaseReg=*/false))
     return false;
 
   return true;
 }
 
 /// Generate an add or subtract for each IVInc in a chain to materialize the IV
 /// user's operand from the previous IV user's operand.
 void LSRInstance::GenerateIVChain(const IVChain &Chain,
if (Best->getType() != Reg->getType() ||
     return false;
   const auto *Diff = dyn_cast<SCEVConstant>(SE.getMinusSCEV(Best, Reg));
   if (!Diff)
     return false;
 
   return TTI.isLegalAddressingMode(
              AccessType.MemTy, /*BaseGV=*/nullptr,
              /*BaseOffset=*/Diff->getAPInt().getSExtValue(),
-             /*HasBaseReg=*/false, /*Scale=*/0, AccessType.AddrSpace) &&
+             /*HasBaseReg=*/true, /*Scale=*/0, AccessType.AddrSpace) &&
          !TTI.isLegalAddressingMode(
              AccessType.MemTy, /*BaseGV=*/nullptr,
              /*BaseOffset=*/-Diff->getAPInt().getSExtValue(),
-             /*HasBaseReg=*/false, /*Scale=*/0, AccessType.AddrSpace);
+             /*HasBaseReg=*/true, /*Scale=*/0, AccessType.AddrSpace);
 }
 
 /// Pick a register which seems likely to be profitable, and then in any use
 /// which has any reference to that register, delete all formulae which do not
 /// reference that register.
 void LSRInstance::NarrowSearchSpaceByPickingWinnerRegs() {
   // With all other options exhausted, loop until the system is simple
   // enough to handle.

llvm/test/Transforms/LoopStrengthReduce/AArch64/lsr-pre-inc-offset-check.ll

 ; preferred addresing mode is preindexed.
 
 
 %"Type" = type <{[166 x [338 x i8]]}>
 define void @test_lsr_pre_inc_offset_check(ptr %p) {
 ; CHECK-LABEL: test_lsr_pre_inc_offset_check:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    mov w8, #165
-; CHECK-NEXT:    add x9, x0, #340
+; CHECK-NEXT:    add x9, x0, #339
 ; CHECK-NEXT:    mov w10, #2
 ; CHECK-NEXT:  .LBB0_1: // %main
 ; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
-; CHECK-NEXT:    stur wzr, [x9, #-1]
+; CHECK-NEXT:    str wzr, [x9]
 ; CHECK-NEXT:    subs x8, x8, #1
-; CHECK-NEXT:    strb w10, [x9]
+; CHECK-NEXT:    strb w10, [x9, #1]
 ; CHECK-NEXT:    add x9, x9, #338
 ; CHECK-NEXT:    b.ne .LBB0_1
 ; CHECK-NEXT:  // %bb.2: // %exit
 ; CHECK-NEXT:    ret
 entry:
   br label %main
 exit:
   ret void

llvm/unittests/Target/AArch64/AddressingModes.cpp

+#include "AArch64Subtarget.h"
+#include "AArch64TargetMachine.h"
+#include "llvm/IR/DataLayout.h"
+#include "llvm/MC/TargetRegistry.h"
+#include "llvm/Support/TargetSelect.h"
+
+#include "gtest/gtest.h"
+#include <initializer_list>
+#include <memory>
+
+using namespace llvm;
+
+namespace {
+
+struct AddrMode : public TargetLowering::AddrMode {
+  constexpr AddrMode(GlobalValue *GV, int64_t Offs, bool HasBase, int64_t S) {
+    BaseGV = GV;
+    BaseOffs = Offs;
+    HasBaseReg = HasBase;
+    Scale = S;
+  }
+};
+struct TestCase {
+  AddrMode AM;
+  unsigned TypeBits;
+  bool Result;
+};
+
+const std::initializer_list<TestCase> Tests = {
+    // {BaseGV, BaseOffs, HasBaseReg, Scale}, Bits, Result
+    {{reinterpret_cast<GlobalValue *>(-1), 0, false, 0}, 64, false},
+    {{nullptr, 8, true, 1}, 64, false},
+    {{nullptr, 0, false, 2}, 64, true},
+    {{nullptr, 0, false, 1}, 64, true},
+    {{nullptr, 4, false, 0}, 64, false},
+
+    {{nullptr, 0, true, 1}, 64, true},
+    {{nullptr, 0, true, 1}, 32, true},
+    {{nullptr, 0, true, 1}, 16, true},
+    {{nullptr, 0, true, 1}, 8, true},
+
+    {{nullptr, 0, true, 2}, 64, false},
+    {{nullptr, 0, true, 2}, 32, false},
+    {{nullptr, 0, true, 2}, 16, true},
+    {{nullptr, 0, true, 2}, 8, false},
+    {{nullptr, 0, true, 4}, 64, false},
+    {{nullptr, 0, true, 4}, 32, true},
+    {{nullptr, 0, true, 4}, 16, false},
+    {{nullptr, 0, true, 4}, 8, false},
+
+    {{nullptr, 0, true, 8}, 64, true},
+    {{nullptr, 0, true, 8}, 32, false},
+    {{nullptr, 0, true, 8}, 16, false},
+    {{nullptr, 0, true, 8}, 8, false},
+
+    {{nullptr, 0, true, 16}, 64, false},
+    {{nullptr, 0, true, 16}, 32, false},
+    {{nullptr, 0, true, 16}, 16, false},
+    {{nullptr, 0, true, 16}, 8, false},
+
+    {{nullptr, -257, true, 0}, 64, false},
+    {{nullptr, -256, true, 0}, 64, true},
+    {{nullptr, -255, true, 0}, 64, true},
+    {{nullptr, -1, true, 0}, 64, true},
+    {{nullptr, 0, true, 0}, 64, true},
+    {{nullptr, 1, true, 0}, 64, true},
+    {{nullptr, 254, true, 0}, 64, true},
+    {{nullptr, 255, true, 0}, 64, true},
+    {{nullptr, 256, true, 0}, 64, true},
+    {{nullptr, 257, true, 0}, 64, false},
+    {{nullptr, 258, true, 0}, 64, false},
+    {{nullptr, 259, true, 0}, 64, false},
+    {{nullptr, 260, true, 0}, 64, false},
+    {{nullptr, 261, true, 0}, 64, false},
+    {{nullptr, 262, true, 0}, 64, false},
+    {{nullptr, 263, true, 0}, 64, false},
+    {{nullptr, 264, true, 0}, 64, true},
+
+    {{nullptr, 4096 * 8 - 8, true, 0}, 64, true},
+    {{nullptr, 4096 * 8 - 7, true, 0}, 64, false},
+    {{nullptr, 4096 * 8 - 6, true, 0}, 64, false},
+    {{nullptr, 4096 * 8 - 5, true, 0}, 64, false},
+    {{nullptr, 4096 * 8 - 4, true, 0}, 64, false},
+    {{nullptr, 4096 * 8 - 3, true, 0}, 64, false},
+    {{nullptr, 4096 * 8 - 2, true, 0}, 64, false},
+    {{nullptr, 4096 * 8 - 1, true, 0}, 64, false},
+    {{nullptr, 4096 * 8, true, 0}, 64, false},
+    {{nullptr, 4096 * 8 + 1, true, 0}, 64, false},
+    {{nullptr, 4096 * 8 + 2, true, 0}, 64, false},
+    {{nullptr, 4096 * 8 + 3, true, 0}, 64, false},
+    {{nullptr, 4096 * 8 + 4, true, 0}, 64, false},
+    {{nullptr, 4096 * 8 + 5, true, 0}, 64, false},
+    {{nullptr, 4096 * 8 + 6, true, 0}, 64, false},
+    {{nullptr, 4096 * 8 + 7, true, 0}, 64, false},
+    {{nullptr, 4096 * 8 + 8, true, 0}, 64, false},
+
+    {{nullptr, -257, true, 0}, 32, false},
+    {{nullptr, -256, true, 0}, 32, true},
+    {{nullptr, -255, true, 0}, 32, true},
+    {{nullptr, -1, true, 0}, 32, true},
+    {{nullptr, 0, true, 0}, 32, true},
+    {{nullptr, 1, true, 0}, 32, true},
+    {{nullptr, 254, true, 0}, 32, true},
+    {{nullptr, 255, true, 0}, 32, true},
+    {{nullptr, 256, true, 0}, 32, true},
+    {{nullptr, 257, true, 0}, 32, false},
+    {{nullptr, 258, true, 0}, 32, false},
+    {{nullptr, 259, true, 0}, 32, false},
+    {{nullptr, 260, true, 0}, 32, true},
+
+    {{nullptr, 4096 * 4 - 4, true, 0}, 32, true},
+    {{nullptr, 4096 * 4 - 3, true, 0}, 32, false},
+    {{nullptr, 4096 * 4 - 2, true, 0}, 32, false},
+    {{nullptr, 4096 * 4 - 1, true, 0}, 32, false},
+    {{nullptr, 4096 * 4, true, 0}, 32, false},
+    {{nullptr, 4096 * 4 + 1, true, 0}, 32, false},
+    {{nullptr, 4096 * 4 + 2, true, 0}, 32, false},
+    {{nullptr, 4096 * 4 + 3, true, 0}, 32, false},
+    {{nullptr, 4096 * 4 + 4, true, 0}, 32, false},
+
+    {{nullptr, -257, true, 0}, 16, false},
+    {{nullptr, -256, true, 0}, 16, true},
+    {{nullptr, -255, true, 0}, 16, true},
+    {{nullptr, -1, true, 0}, 16, true},
+    {{nullptr, 0, true, 0}, 16, true},
+    {{nullptr, 1, true, 0}, 16, true},
+    {{nullptr, 254, true, 0}, 16, true},
+    {{nullptr, 255, true, 0}, 16, true},
+    {{nullptr, 256, true, 0}, 16, true},
+    {{nullptr, 257, true, 0}, 16, false},
+    {{nullptr, 258, true, 0}, 16, true},
+
+    {{nullptr, 4096 * 2 - 2, true, 0}, 16, true},
+    {{nullptr, 4096 * 2 - 1, true, 0}, 16, false},
+    {{nullptr, 4096 * 2, true, 0}, 16, false},
+    {{nullptr, 4096 * 2 + 1, true, 0}, 16, false},
+    {{nullptr, 4096 * 2 + 2, true, 0}, 16, false},
+
+    {{nullptr, -257, true, 0}, 8, false},
+    {{nullptr, -256, true, 0}, 8, true},
+    {{nullptr, -255, true, 0}, 8, true},
+    {{nullptr, -1, true, 0}, 8, true},
+    {{nullptr, 0, true, 0}, 8, true},
+    {{nullptr, 1, true, 0}, 8, true},
+    {{nullptr, 254, true, 0}, 8, true},
+    {{nullptr, 255, true, 0}, 8, true},
+    {{nullptr, 256, true, 0}, 8, true},
+    {{nullptr, 257, true, 0}, 8, true},
+
+    {{nullptr, 4096 - 2, true, 0}, 8, true},
+    {{nullptr, 4096 - 1, true, 0}, 8, true},
+    {{nullptr, 4096, true, 0}, 8, false},
+    {{nullptr, 4096 + 1, true, 0}, 8, false},
+
+};
+} // namespace
+
+TEST(AddressingModes, AddressingModes) {
+  LLVMInitializeAArch64TargetInfo();
+  LLVMInitializeAArch64Target();
+  LLVMInitializeAArch64TargetMC();
+
+  std::string Error;
+  auto TT = Triple::normalize("aarch64");
+  const Target *T = TargetRegistry::lookupTarget(TT, Error);
+
+  std::unique_ptr<TargetMachine> TM(
+      T->createTargetMachine(TT, "generic", "", TargetOptions(), std::nullopt,
+                             std::nullopt, CodeGenOpt::Default));
+  AArch64Subtarget ST(TM->getTargetTriple(), TM->getTargetCPU(),
+                      TM->getTargetCPU(), TM->getTargetFeatureString(), *TM,
+                      true);
+
+  auto *TLI = ST.getTargetLowering();
+  DataLayout DL("e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128");
+  LLVMContext Ctx;
+
+  for (const auto &Test : Tests) {
+    Type *Typ = Type::getIntNTy(Ctx, Test.TypeBits);
+    ASSERT_EQ(TLI->isLegalAddressingMode(DL, Test.AM, Typ, 0), Test.Result);
+  }
+}

llvm/unittests/Target/AArch64/CMakeLists.txt

set(LLVM_LINK_COMPONENTS
   SelectionDAG
   Support
   Target
   TargetParser
   )
 
 add_llvm_target_unittest(AArch64Tests
   AArch64InstPrinterTest.cpp
+  AddressingModes.cpp
   DecomposeStackOffsetTest.cpp
   InstSizes.cpp
   MatrixRegisterAliasing.cpp
   SMEAttributesTest.cpp
   )
 
 set_property(TARGET AArch64Tests PROPERTY FOLDER "Tests/UnitTests/TargetTests")