[llvm] [AArch64] merge index address with large offset into base address (PR #72187)

via llvm-commits llvm-commits at lists.llvm.org
Wed Nov 29 01:28:38 PST 2023


https://github.com/vfdff updated https://github.com/llvm/llvm-project/pull/72187

>From 42e4cf5095025c4f7ec4f853ee2a060d92c88af7 Mon Sep 17 00:00:00 2001
From: zhongyunde 00443407 <zhongyunde at huawei.com>
Date: Mon, 20 Nov 2023 01:13:43 -0500
Subject: [PATCH 1/2] [AArch64] merge index address with large offset into base
 address

A case for this transformation, https://gcc.godbolt.org/z/nhYcWq1WE
```
Fold
  mov     w8, #56952
  movk    w8, #15, lsl #16
  ldrb    w0, [x0, x8]
into
  add     x0, x0, 1036288
  ldrb    w0, [x0, 3704]
```
Only support single use base for the first time.
Fix https://github.com/llvm/llvm-project/issues/71917

TODO: support the multiple-uses with reuseing common base offset.
https://gcc.godbolt.org/z/Mr7srTjnz
---
 .../Target/AArch64/AArch64ISelDAGToDAG.cpp    | 64 ++++++++++++++
 llvm/test/CodeGen/AArch64/arm64-addrmode.ll   | 86 +++++++++++++++++++
 2 files changed, 150 insertions(+)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
index 7617dccdeee397f..f67dc1f5156ec41 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
@@ -1069,6 +1069,42 @@ bool AArch64DAGToDAGISel::SelectAddrModeIndexedBitWidth(SDValue N, bool IsSigned
   return true;
 }
 
+// 16-bit optionally shifted immediates are legal for single mov.
+static bool isLegalSingleMOVImmediate(int64_t Immed) {
+  if (Immed == std::numeric_limits<int64_t>::min()) {
+    LLVM_DEBUG(dbgs() << "Illegal single mov imm " << Immed
+                      << ": avoid UB for INT64_MIN\n");
+    return false;
+  }
+  // The shift value can be 0(default), 16, 32, 48
+  bool IsLegal =
+      ((Immed >> 16) == 0 || ((Immed & 0xffff) == 0 && Immed >> 32 == 0) ||
+       ((Immed & 0xffffffff) == 0 && Immed >> 48 == 0) ||
+       ((Immed & 0xffffffffffff) == 0));
+  LLVM_DEBUG(dbgs() << "Is " << Immed << " legal single mov imm: "
+                    << (IsLegal ? "yes" : "no") << "\n");
+  return IsLegal;
+}
+
+// Check whether a unsigned vaule is not in the immediate range of mov but in
+// the immediate range of imm24. The "Size" argument is the size in bytes of the
+// memory reference.
+static bool isPreferredBaseAddrMode(int64_t ImmOff, unsigned Size) {
+  if ((ImmOff & (Size - 1)) != 0 || ImmOff < 0)
+    return false;
+
+  // If the immediate already can be encoded in mov, then just keep the existing
+  // logic.
+  if (isLegalSingleMOVImmediate(ImmOff))
+    return false;
+
+  // For a imm24, its low imm12 can be fold as the immediate of load or store,
+  // and its high part can be encoded in an add.
+  if (ImmOff >> 24 == 0)
+    return true;
+  return false;
+}
+
 /// SelectAddrModeIndexed - Select a "register plus scaled unsigned 12-bit
 /// immediate" address.  The "Size" argument is the size in bytes of the memory
 /// reference, which determines the scale.
@@ -1110,6 +1146,25 @@ bool AArch64DAGToDAGISel::SelectAddrModeIndexed(SDValue N, unsigned Size,
         OffImm = CurDAG->getTargetConstant(RHSC >> Scale, dl, MVT::i64);
         return true;
       }
+
+      // Perfer [Reg + imm] mode.
+      //     ADD  BaseReg, WideImmediate & 0x0fff000
+      //     LDR  X2, [BaseReg, WideImmediate & 0x0fff]
+      // TODO: should reuse the base add for multiple-uses.
+      SDValue LHS = N.getOperand(0);
+      if (N.hasOneUse() && isPreferredBaseAddrMode(RHSC, Size)) {
+        int64_t ImmOffUnScale = RHSC;
+        int64_t ImmOffLow = ImmOffUnScale & 0x0fff;
+        int64_t ImmOffHigh = RHSC - ImmOffLow;
+        SDValue ImmHighSDV =
+            CurDAG->getTargetConstant(ImmOffHigh >> 12, dl, MVT::i64);
+        Base = SDValue(CurDAG->getMachineNode(
+                           AArch64::ADDXri, dl, MVT::i64, LHS, ImmHighSDV,
+                           CurDAG->getTargetConstant(12, dl, MVT::i32)),
+                       0);
+        OffImm = CurDAG->getTargetConstant(ImmOffLow >> Scale, dl, MVT::i64);
+        return true;
+      }
     }
   }
 
@@ -1293,6 +1348,8 @@ bool AArch64DAGToDAGISel::SelectAddrModeXRO(SDValue N, unsigned Size,
     return false;
   SDValue LHS = N.getOperand(0);
   SDValue RHS = N.getOperand(1);
+  // N may be update later, so get its status here.
+  bool HasOneUse = N.hasOneUse();
   SDLoc DL(N);
 
   // Check if this particular node is reused in any non-memory related
@@ -1351,6 +1408,13 @@ bool AArch64DAGToDAGISel::SelectAddrModeXRO(SDValue N, unsigned Size,
     return true;
   }
 
+  // Perfer [Reg + imm] mode, so skip this scenarios.
+  if (auto *OffsetC = dyn_cast<ConstantSDNode>(RHS)) {
+    int64_t ImmOff = (int64_t)OffsetC->getZExtValue();
+    if (HasOneUse && isPreferredBaseAddrMode(ImmOff, Size)) {
+      return false;
+    }
+  }
   // Match any non-shifted, non-extend, non-immediate add expression.
   Base = LHS;
   Offset = RHS;
diff --git a/llvm/test/CodeGen/AArch64/arm64-addrmode.ll b/llvm/test/CodeGen/AArch64/arm64-addrmode.ll
index 69c558d9d5599dc..d3543992e4598b3 100644
--- a/llvm/test/CodeGen/AArch64/arm64-addrmode.ll
+++ b/llvm/test/CodeGen/AArch64/arm64-addrmode.ll
@@ -209,3 +209,89 @@ define void @t17(i64 %a) {
   %3 = load volatile i64, ptr %2, align 8
   ret void
 }
+
+; https://gcc.godbolt.org/z/ErhhdxMv3
+define i32 @LdOffset_i8(ptr %a)  {
+; CHECK-LABEL: LdOffset_i8:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add x8, x0, #253, lsl #12 // =1036288
+; CHECK-NEXT:    ldrb w0, [x8, #3704]
+; CHECK-NEXT:    ret
+  %arrayidx = getelementptr inbounds i8, ptr %a, i64 1039992
+  %val = load i8, ptr %arrayidx, align 1
+  %conv = zext i8 %val to i32
+  ret i32 %conv
+}
+
+define i32 @LdOffset_i16(ptr %a)  {
+; CHECK-LABEL: LdOffset_i16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add x8, x0, #507, lsl #12 // =2076672
+; CHECK-NEXT:    ldrsh w0, [x8, #3312]
+; CHECK-NEXT:    ret
+  %arrayidx = getelementptr inbounds i16, ptr %a, i64 1039992
+  %val = load i16, ptr %arrayidx, align 2
+  %conv = sext i16 %val to i32
+  ret i32 %conv
+}
+
+define i32 @LdOffset_i32(ptr %a)  {
+; CHECK-LABEL: LdOffset_i32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add x8, x0, #1015, lsl #12 // =4157440
+; CHECK-NEXT:    ldr w0, [x8, #2528]
+; CHECK-NEXT:    ret
+  %arrayidx = getelementptr inbounds i32, ptr %a, i64 1039992
+  %val = load i32, ptr %arrayidx, align 4
+  ret i32 %val
+}
+
+; TODO: https://gcc.godbolt.org/z/q6frE9ePe
+define i64 @LdOffset_i64_multi_offset(ptr %a) {
+; CHECK-LABEL: LdOffset_i64_multi_offset:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add x8, x0, #2031, lsl #12 // =8318976
+; CHECK-NEXT:    add x8, x8, #960
+; CHECK-NEXT:    ldr x9, [x8]
+; CHECK-NEXT:    ldr x8, [x8, #2056]
+; CHECK-NEXT:    add x0, x8, x9
+; CHECK-NEXT:    ret
+  %arrayidx = getelementptr inbounds i64, ptr %a, i64 1039992
+  %val0 = load i64, ptr %arrayidx, align 8
+  %arrayidx1 = getelementptr inbounds i64, ptr %a, i64 1040249
+  %val1 = load i64, ptr %arrayidx1, align 8
+  %add = add nsw i64 %val1, %val0
+  ret i64 %add
+}
+
+define i64 @LdOffset_i64_multi_offset_with_commmon_base(ptr %a) {
+; CHECK-LABEL: LdOffset_i64_multi_offset_with_commmon_base:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add x8, x0, #507, lsl #12 // =2076672
+; CHECK-NEXT:    ldr x9, [x8, #26464]
+; CHECK-NEXT:    ldr x8, [x8, #26496]
+; CHECK-NEXT:    add x0, x8, x9
+; CHECK-NEXT:    ret
+  %b = getelementptr inbounds i16, ptr %a, i64 1038336
+  %arrayidx = getelementptr inbounds i64, ptr %b, i64 3308
+  %val0 = load i64, ptr %arrayidx, align 8
+  %arrayidx1 = getelementptr inbounds i64, ptr %b, i64 3312
+  %val1 = load i64, ptr %arrayidx1, align 8
+  %add = add nsw i64 %val1, %val0
+  ret i64 %add
+}
+
+; Negative test: the offset is odd
+define i32 @LdOffset_i16_odd_offset(ptr nocapture noundef readonly %a)  {
+; CHECK-LABEL: LdOffset_i16_odd_offset:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w8, #56953 // =0xde79
+; CHECK-NEXT:    movk w8, #15, lsl #16
+; CHECK-NEXT:    ldrsh w0, [x0, x8]
+; CHECK-NEXT:    ret
+  %arrayidx = getelementptr inbounds i8, ptr %a, i64 1039993
+  %val = load i16, ptr %arrayidx, align 2
+  %conv = sext i16 %val to i32
+  ret i32 %conv
+}
+

>From 49e5f3a138c006b5177c1dc98e5827c32225cdfb Mon Sep 17 00:00:00 2001
From: zhongyunde 00443407 <zhongyunde at huawei.com>
Date: Tue, 28 Nov 2023 07:05:51 -0500
Subject: [PATCH 2/2] [AArch64] Rebase the common base offset for better ISel

When all the large const offsets masked with the same value from bit-12 to bit-23.
Fold
  add     x8, x0, #2031, lsl #12
  add     x8, x8, #960
  ldr     x9, [x8, x8]
  ldr     x8, [x8, #2056]

into
  add     x8, x0, #2031, lsl #12
  ldr     x9, [x8, #960]
  ldr     x8, [x8, #3016]
---
 llvm/include/llvm/CodeGen/BasicTTIImpl.h      |  4 ++++
 llvm/include/llvm/CodeGen/TargetLowering.h    |  6 +++++
 llvm/lib/CodeGen/CodeGenPrepare.cpp           | 23 +++++++++++++++++++
 .../Target/AArch64/AArch64ISelDAGToDAG.cpp    |  6 ++---
 .../Target/AArch64/AArch64ISelLowering.cpp    | 12 ++++++++++
 llvm/lib/Target/AArch64/AArch64ISelLowering.h |  3 +++
 llvm/test/CodeGen/AArch64/arm64-addrmode.ll   |  7 +++---
 7 files changed, 53 insertions(+), 8 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/BasicTTIImpl.h b/llvm/include/llvm/CodeGen/BasicTTIImpl.h
index 7a8f36da58ceccb..e367705630c2b4d 100644
--- a/llvm/include/llvm/CodeGen/BasicTTIImpl.h
+++ b/llvm/include/llvm/CodeGen/BasicTTIImpl.h
@@ -342,6 +342,10 @@ class BasicTTIImplBase : public TargetTransformInfoImplCRTPBase<T> {
     return getTLI()->isLegalAddressingMode(DL, AM, Ty, AddrSpace, I);
   }
 
+  int64_t getPreferBaseOffset(int64_t MinOffset, int64_t MaxOffset) {
+    return getTLI()->getPreferBaseOffset(MinOffset, MaxOffset);
+  }
+
   unsigned getStoreMinimumVF(unsigned VF, Type *ScalarMemTy,
                              Type *ScalarValTy) const {
     auto &&IsSupportedByTarget = [this, ScalarMemTy, ScalarValTy](unsigned VF) {
diff --git a/llvm/include/llvm/CodeGen/TargetLowering.h b/llvm/include/llvm/CodeGen/TargetLowering.h
index 9ebcc28c38ae672..ecaaa2cf9f7a1a8 100644
--- a/llvm/include/llvm/CodeGen/TargetLowering.h
+++ b/llvm/include/llvm/CodeGen/TargetLowering.h
@@ -2720,6 +2720,12 @@ class TargetLoweringBase {
                                      Type *Ty, unsigned AddrSpace,
                                      Instruction *I = nullptr) const;
 
+  /// Return the perfered common base offset.
+  virtual int64_t getPreferBaseOffset(int64_t MinOffset,
+                                      int64_t MaxOffset) const {
+    return 0;
+  }
+
   /// Return true if the specified immediate is legal icmp immediate, that is
   /// the target has icmp instructions which can compare a register against the
   /// immediate without having to materialize the immediate into a register.
diff --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp
index 07dc718ee3a38c3..818c40e0530c635 100644
--- a/llvm/lib/CodeGen/CodeGenPrepare.cpp
+++ b/llvm/lib/CodeGen/CodeGenPrepare.cpp
@@ -6099,6 +6099,29 @@ bool CodeGenPrepare::splitLargeGEPOffsets() {
     int64_t BaseOffset = LargeOffsetGEPs.begin()->second;
     Value *NewBaseGEP = nullptr;
 
+    // Check whether all the offsets can be encoded with perfered common base.
+    if (int64_t PreferBase = TLI->getPreferBaseOffset(
+            LargeOffsetGEPs.front().second, LargeOffsetGEPs.back().second)) {
+      BaseOffset = PreferBase;
+
+      // Create a new base.
+      LLVMContext &Ctx = BaseGEP->getContext();
+      Type *PtrIdxTy = DL->getIndexType(BaseGEP->getType());
+      Type *I8PtrTy =
+          PointerType::get(Ctx, BaseGEP->getType()->getPointerAddressSpace());
+      Type *I8Ty = Type::getInt8Ty(Ctx);
+      BasicBlock *NewBaseInsertBB = &BaseGEP->getFunction()->getEntryBlock();
+      BasicBlock::iterator NewBaseInsertPt =
+          NewBaseInsertBB->getFirstInsertionPt();
+      IRBuilder<> NewBaseBuilder(NewBaseInsertBB, NewBaseInsertPt);
+
+      Value *BaseIndex = ConstantInt::get(PtrIdxTy, BaseOffset);
+      NewBaseGEP =
+          NewBaseBuilder.CreateGEP(I8Ty, OldBase, BaseIndex, "splitgep");
+      LLVM_DEBUG(dbgs() << "CGP: New common base: " << *NewBaseGEP << "\n");
+      NewGEPBases.insert(NewBaseGEP);
+    }
+
     auto *LargeOffsetGEP = LargeOffsetGEPs.begin();
     while (LargeOffsetGEP != LargeOffsetGEPs.end()) {
       GetElementPtrInst *GEP = LargeOffsetGEP->first;
diff --git a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
index f67dc1f5156ec41..77a626fe0c17edc 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
@@ -1152,7 +1152,7 @@ bool AArch64DAGToDAGISel::SelectAddrModeIndexed(SDValue N, unsigned Size,
       //     LDR  X2, [BaseReg, WideImmediate & 0x0fff]
       // TODO: should reuse the base add for multiple-uses.
       SDValue LHS = N.getOperand(0);
-      if (N.hasOneUse() && isPreferredBaseAddrMode(RHSC, Size)) {
+      if (isPreferredBaseAddrMode(RHSC, Size)) {
         int64_t ImmOffUnScale = RHSC;
         int64_t ImmOffLow = ImmOffUnScale & 0x0fff;
         int64_t ImmOffHigh = RHSC - ImmOffLow;
@@ -1348,8 +1348,6 @@ bool AArch64DAGToDAGISel::SelectAddrModeXRO(SDValue N, unsigned Size,
     return false;
   SDValue LHS = N.getOperand(0);
   SDValue RHS = N.getOperand(1);
-  // N may be update later, so get its status here.
-  bool HasOneUse = N.hasOneUse();
   SDLoc DL(N);
 
   // Check if this particular node is reused in any non-memory related
@@ -1411,7 +1409,7 @@ bool AArch64DAGToDAGISel::SelectAddrModeXRO(SDValue N, unsigned Size,
   // Perfer [Reg + imm] mode, so skip this scenarios.
   if (auto *OffsetC = dyn_cast<ConstantSDNode>(RHS)) {
     int64_t ImmOff = (int64_t)OffsetC->getZExtValue();
-    if (HasOneUse && isPreferredBaseAddrMode(ImmOff, Size)) {
+    if (isPreferredBaseAddrMode(ImmOff, Size)) {
       return false;
     }
   }
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index d42ae4ff93a4442..e62f420258bd367 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -15977,6 +15977,18 @@ bool AArch64TargetLowering::isLegalAddressingMode(const DataLayout &DL,
                                                           AM.Scale);
 }
 
+// Check whether the 2 offsets belong to the same imm24 range, and their high
+// 12bits are same.
+int64_t AArch64TargetLowering::getPreferBaseOffset(int64_t MinOffset,
+                                                   int64_t MaxOffset) const {
+  if (MinOffset >> 12 == MaxOffset >> 12 && MinOffset >> 24 == 0) {
+    // Rebase the value to an integer multiple of imm12.
+    return MinOffset & 0xfff000;
+  }
+
+  return 0;
+}
+
 bool AArch64TargetLowering::shouldConsiderGEPOffsetSplit() const {
   // Consider splitting large offset of struct or array.
   return true;
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.h b/llvm/lib/Target/AArch64/AArch64ISelLowering.h
index 2a039488f2a9ab3..a7d731abd0cb575 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.h
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.h
@@ -679,6 +679,9 @@ class AArch64TargetLowering : public TargetLowering {
                              unsigned AS,
                              Instruction *I = nullptr) const override;
 
+  int64_t getPreferBaseOffset(int64_t MinOffset,
+                              int64_t MaxOffset) const override;
+
   /// Return true if an FMA operation is faster than a pair of fmul and fadd
   /// instructions. fmuladd intrinsics will be expanded to FMAs when this method
   /// returns true, otherwise fmuladd is expanded to fmul + fadd.
diff --git a/llvm/test/CodeGen/AArch64/arm64-addrmode.ll b/llvm/test/CodeGen/AArch64/arm64-addrmode.ll
index d3543992e4598b3..1497396a0dcd1d3 100644
--- a/llvm/test/CodeGen/AArch64/arm64-addrmode.ll
+++ b/llvm/test/CodeGen/AArch64/arm64-addrmode.ll
@@ -246,14 +246,13 @@ define i32 @LdOffset_i32(ptr %a)  {
   ret i32 %val
 }
 
-; TODO: https://gcc.godbolt.org/z/q6frE9ePe
+; https://gcc.godbolt.org/z/q6frE9ePe
 define i64 @LdOffset_i64_multi_offset(ptr %a) {
 ; CHECK-LABEL: LdOffset_i64_multi_offset:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    add x8, x0, #2031, lsl #12 // =8318976
-; CHECK-NEXT:    add x8, x8, #960
-; CHECK-NEXT:    ldr x9, [x8]
-; CHECK-NEXT:    ldr x8, [x8, #2056]
+; CHECK-NEXT:    ldr x9, [x8, #960]
+; CHECK-NEXT:    ldr x8, [x8, #3016]
 ; CHECK-NEXT:    add x0, x8, x9
 ; CHECK-NEXT:    ret
   %arrayidx = getelementptr inbounds i64, ptr %a, i64 1039992



More information about the llvm-commits mailing list