[AArch64][ISel] Always use pre-inc/post-inc addressing mode for auto-indexed load/store with constant offset.

[huihuiz]

Summary

Unlike ARM target, current AArch64 target doesn't have facility to encode the
operation bit: whether to add an offset to base pointer for pre-inc/post-inc
addressing mode, or to subtract an offset from base pointer for
pre-dec/post-dec addressing mode.

A mis-compile (https://github.com/llvm/llvm-project/issues/60645) was noticed
due to this limitation.

Therefore, for AArch64 auto-indexed load/store with constant offset, always
use pre-inc/post-inc addressing mode. The constant offset is negated for
pre-dec/post-dec addressing mode.
An auto-indexed address with non-constant offset is currently not split into
base and offset parts. If we are to handle non-constant offset in the future,
offset node will need to take a negate.

[huihuiz]

Test case was reduced from internal benchmark. Assuming we would run into similar issues for post-dec addressing mode.
If needed I can try to synthesize a test for post-dec.

[efriedma]

The way this fix is implemented seems suspicious. As far as I can tell, this would basically imply that PRE_DEC never worked on any target. A quick grep shows three targets using it. If the other targets are working, we should probably be fixing this in AArch64-specific code.

Apparently, the AArch64 backend can also generate PRE_INC operations with a negative offset, which might explain why nobody noticed this issue before on AArch64.

[huihuiz]

Dumped -debug-only=isel,dagcombine for a few targets (armv7a, aarch64, x86_64, riscv64) to see what's going on. Here is what I found:

1, Not all targets support PRE_DEC addressing mode, we only replace a store with equivalent <pre-dec> store when it's legal to do so.
Take previous AArch64 t.ll for example

t6: i64 = add nuw nsw t4, Constant:i64<8>
t15: i64 = shl t6, Constant:i64<2>
t23: i64 = sub t2, t15
t13: ch = store<(store (s32) into %ir.t4)> t0, Constant:i32<0>, t23, undef:i64

t13 was replaced with a <pre-dec> store t32

t27: i64 = shl t4, Constant:i64<2>
t30: i64 = sub t2, t27
t32: i64,ch = store<(store (s32) into %ir.t4), <pre-dec>> t0, Constant:i32<0>, t30, Constant:i64<32>

In DAGCombiner::CombineToPreIndexedLoadStore(), before DAG.getIndexedStore() is called, we first check TLI.isIndexedStoreLegal() in getCombineLoadStoreParts().
Then we set AddrMode AM properly.

For targets like x86_64, index store is not legal.

2, For armv7a target, I change the attached t.ll to use 32-bit as gep offset. CodeGen is correct, for two reasons:
i) For ARM target lowering, commute-with-shift is turned off after legalization, but for AArch64 it's returning true when checking isDesiableToCommuteWithShift(). So constant 32 is pulled out for AArch64 target, but not for ARM target.
ii) Negating offset register for ARM target is done in SelectAddrMode2OffsetReg() , the ARM_AM::getAM2Opc() part.
We generate:
add r1, r1, #8
str r2, [r3, -r1, lsl #2]! // correct

When Selecting node t26
t6: i32 = add nuw nsw t4, Constant:i32<8>
t14: i32 = shl t6, Constant:i32<2>
t26: i32,ch = store<(store (s32) into %ir.t4), <pre-dec>> t0, Constant:i32<0>, t2, t14

we got
Morphed node: t26: i32,ch = STRr_preidx<Mem:(store (s32) into %ir.t4)> Constant:i32<0>, t2, t6, TargetConstant:i32<20482>, TargetConstant:i32<14>, Register:i32 $noreg, t0

TargetConstant:i32<20482> , bit 13 is set to 1 to indicate doing subtraction.

3, I don't seem to find similar handling for AArch64 target. If I am looking at the wrong direction, please point me out?

[efriedma]

I think the targets that use PRE_DEC are ARM and AVR. If existing testcases for PRE_DEC work on those targets, maybe we don't need to be too concerned about other targets? I guess maybe they don't hit this code, somehow?

This still doesn't seem conceptually correct, though. A PRE_DEC addressing mode doesn't have to be an immediate; it can be an arbitrary runtime value. And there isn't any way to negate an arbitrary value without creating extra nodes. So I'd expect any necessary negation to happen in target-specific code.

Probably the simplest thing is just to fix the AArch64 code so it never generates PRE_DEC/POST_DEC. They really just make things more complicated, and we don't need them on AArch64 because the all PRE_INC/PRE_DEC ops have an immediate offset.

The pre_store pattern in TableGen seems like sort of a trap: it encompasses both PRE_INC and PRE_DEC, but doesn't provide any way to actually tell whether you matched a PRE_INC or a PRE_DEC.

[huihuiz]

Thanks Eli for the feedbacks!

Current AArch64TargetLowering::getIndexedAddressParts() is not splitting out a non-constant offset. If we are to split out a non-constant offset, we need to make sure a negate is applied to that offset.

take another t.ll below , using non-constant offset

define i8* @test(i8* %ptr, i64 %t0, i64 %off) {
  %t1 = add nuw nsw i64 %t0, %off
  %t2 = mul i64 %t1, -4
  %t3 = getelementptr i8, i8* %ptr, i64 %t2
  %t4 = bitcast i8* %t3 to i32*
  store i32 0, i32* %t4, align 4
  %t5 = shl i64 %t1, 2
  %t6 = sub nuw nsw i64 -8, %t5
  %t7 = getelementptr i8, i8* %ptr, i64 %t6
  %t8 = bitcast i8* %t7 to i32*
  store i32 0, i32* %t8, align 4
  ret i8* %ptr
}

We are generating

// %bb.0:
        add     x8, x1, x2
        sub     x8, x0, x8, lsl #2
        str     wzr, [x8]  // We are not splitting out a non-constant offset
        stur    wzr, [x8, #-8]
        ret

[efriedma]

The logic here seems fine, but please clean up the dead code (the IsInc argument of getIndexedAddressParts is always set to true, so you can just remove it and clean up the callers).

[huihuiz]

Addressed review comments, also clean out "ISD::MemIndexedMode &AM" from getIndexedAddressParts().

[efriedma]

LGTM

llvm/lib/Target/AArch64/AArch64ISelLowering.h

   }
 
   bool shouldExtendGSIndex(EVT VT, EVT &EltTy) const override;
   bool shouldRemoveExtendFromGSIndex(EVT IndexVT, EVT DataVT) const override;
   bool isVectorLoadExtDesirable(SDValue ExtVal) const override;
   bool isUsedByReturnOnly(SDNode *N, SDValue &Chain) const override;
   bool mayBeEmittedAsTailCall(const CallInst *CI) const override;
   bool getIndexedAddressParts(SDNode *N, SDNode *Op, SDValue &Base,
-                              SDValue &Offset, ISD::MemIndexedMode &AM,
-                              bool &IsInc, SelectionDAG &DAG) const;
+                              SDValue &Offset, SelectionDAG &DAG) const;
   bool getPreIndexedAddressParts(SDNode *N, SDValue &Base, SDValue &Offset,
                                  ISD::MemIndexedMode &AM,
                                  SelectionDAG &DAG) const override;
   bool getPostIndexedAddressParts(SDNode *N, SDNode *Op, SDValue &Base,
                                   SDValue &Offset, ISD::MemIndexedMode &AM,
                                   SelectionDAG &DAG) const override;
 
   void ReplaceNodeResults(SDNode *N, SmallVectorImpl<SDValue> &Results,

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

 // Return whether the an instruction can potentially be optimized to a tail
 // call. This will cause the optimizers to attempt to move, or duplicate,
 // return instructions to help enable tail call optimizations for this
 // instruction.
 bool AArch64TargetLowering::mayBeEmittedAsTailCall(const CallInst *CI) const {
   return CI->isTailCall();
 }
 
-bool AArch64TargetLowering::getIndexedAddressParts(
-    SDNode *N, SDNode *Op, SDValue &Base, SDValue &Offset,
-    ISD::MemIndexedMode &AM, bool &IsInc, SelectionDAG &DAG) const {
+bool AArch64TargetLowering::getIndexedAddressParts(SDNode *N, SDNode *Op,
+                                                   SDValue &Base,
+                                                   SDValue &Offset,
+                                                   SelectionDAG &DAG) const {
   if (Op->getOpcode() != ISD::ADD && Op->getOpcode() != ISD::SUB)
     return false;
 
   // Non-null if there is exactly one user of the loaded value (ignoring chain).
   SDNode *ValOnlyUser = nullptr;
   for (SDNode::use_iterator UI = N->use_begin(), UE = N->use_end(); UI != UE;
        ++UI) {
     if (UI.getUse().getResNo() == 1)

   // All of the indexed addressing mode instructions take a signed
   // 9 bit immediate offset.
   if (ConstantSDNode *RHS = dyn_cast<ConstantSDNode>(Op->getOperand(1))) {
     int64_t RHSC = RHS->getSExtValue();
     if (Op->getOpcode() == ISD::SUB)
       RHSC = -(uint64_t)RHSC;
     if (!isInt<9>(RHSC))
       return false;
-    IsInc = (Op->getOpcode() == ISD::ADD);
-    Offset = Op->getOperand(1);
+    // Always emit pre-inc/post-inc addressing mode. Use negated constant offset
+    // when dealing with subtraction.
+    Offset = DAG.getConstant(RHSC, SDLoc(N), RHS->getValueType(0));
     return true;
   }
   return false;
 }
 
 bool AArch64TargetLowering::getPreIndexedAddressParts(SDNode *N, SDValue &Base,
                                                       SDValue &Offset,
                                                       ISD::MemIndexedMode &AM,
                                                       SelectionDAG &DAG) const {
   EVT VT;
   SDValue Ptr;
   if (LoadSDNode *LD = dyn_cast<LoadSDNode>(N)) {
     VT = LD->getMemoryVT();
     Ptr = LD->getBasePtr();
   } else if (StoreSDNode *ST = dyn_cast<StoreSDNode>(N)) {
     VT = ST->getMemoryVT();
     Ptr = ST->getBasePtr();
   } else
     return false;
 
-  bool IsInc;
-  if (!getIndexedAddressParts(N, Ptr.getNode(), Base, Offset, AM, IsInc, DAG))
+  if (!getIndexedAddressParts(N, Ptr.getNode(), Base, Offset, DAG))
     return false;
-  AM = IsInc ? ISD::PRE_INC : ISD::PRE_DEC;
+  AM = ISD::PRE_INC;
   return true;
 }
 
 bool AArch64TargetLowering::getPostIndexedAddressParts(
     SDNode *N, SDNode *Op, SDValue &Base, SDValue &Offset,
     ISD::MemIndexedMode &AM, SelectionDAG &DAG) const {
   EVT VT;
   SDValue Ptr;
   if (LoadSDNode *LD = dyn_cast<LoadSDNode>(N)) {
     VT = LD->getMemoryVT();
     Ptr = LD->getBasePtr();
   } else if (StoreSDNode *ST = dyn_cast<StoreSDNode>(N)) {
     VT = ST->getMemoryVT();
     Ptr = ST->getBasePtr();
   } else
     return false;
 
-  bool IsInc;
-  if (!getIndexedAddressParts(N, Op, Base, Offset, AM, IsInc, DAG))
+  if (!getIndexedAddressParts(N, Op, Base, Offset, DAG))
     return false;
   // Post-indexing updates the base, so it's not a valid transform
   // if that's not the same as the load's pointer.
   if (Ptr != Base)
     return false;
-  AM = IsInc ? ISD::POST_INC : ISD::POST_DEC;
+  AM = ISD::POST_INC;
   return true;
 }
 
 void AArch64TargetLowering::ReplaceBITCASTResults(
     SDNode *N, SmallVectorImpl<SDValue> &Results, SelectionDAG &DAG) const {
   SDLoc DL(N);
   SDValue Op = N->getOperand(0);
   EVT VT = N->getValueType(0);

llvm/test/CodeGen/AArch64/pre-indexed-addrmode-with-constant-offset.ll

+; RUN: llc -mtriple=aarch64-linux-gnu < %s | FileCheck %s
+
+; Reduced test from https://github.com/llvm/llvm-project/issues/60645.
+; To check that we are generating -32 as offset for the first store.
+
+define i8* @pr60645(i8* %ptr, i64 %t0) {
+; CHECK-LABEL: pr60645:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    sub x8, x0, x1, lsl #2
+; CHECK-NEXT:    str wzr, [x8, #-32]!
+; CHECK-NEXT:    stur wzr, [x8, #-8]
+; CHECK-NEXT:    ret
+  %t1 = add nuw nsw i64 %t0, 8
+  %t2 = mul i64 %t1, -4
+  %t3 = getelementptr i8, i8* %ptr, i64 %t2
+  %t4 = bitcast i8* %t3 to i32*
+  store i32 0, i32* %t4, align 4
+  %t5 = shl i64 %t1, 2
+  %t6 = sub nuw nsw i64 -8, %t5
+  %t7 = getelementptr i8, i8* %ptr, i64 %t6
+  %t8 = bitcast i8* %t7 to i32*
+  store i32 0, i32* %t8, align 4
+  ret i8* %ptr
+}