[DAG] Teach BaseIndexOffset to correctly handle with indexed operations

[niravd]

Summary

BaseIndexOffset address analysis incorrectly ignores offsets folded
into indexed memory operations causing potential errors in alias
analysis of pre-indexed operations.

[efriedma]

I would prefer to have a regression test, but it's okay if you can't come up with a way to write one.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGAddressAnalysis.cpp

90

Does this need to be something like "BaseIndexOffset(Base, N->getOffset(), 0, false);"? Currently, it looks like this code is discarding the increment.

[niravd]

Correct failing on non-constant indices.

[niravd]

In D41701#966881, @efriedma wrote:

I would prefer to have a regression test, but it's okay if you can't come up with a way to write one.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGAddressAnalysis.cpp

90

Good catch. This won't work on pre-decrement though as we'd need the negated index. I've changed it to mark such cases as returning no information and modifying other uses to check for a null base.

[efriedma]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

13360

Does this null check still work properly given that BaseIndexOffset::match can return a null base?

[niravd]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

13360

The indexed check @ line 13350 precludes the new case from happening here.

[efriedma]

LGTM

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

13360

Oh, I see.

llvm/include/llvm/CodeGen/SelectionDAGAddressAnalysis.h

bool equalBaseIndex(BaseIndexOffset &Other, const SelectionDAG &DAG) {
     int64_t Off;
     return equalBaseIndex(Other, DAG, Off);
   }
 
   bool equalBaseIndex(BaseIndexOffset &Other, const SelectionDAG &DAG,
                       int64_t &Off);
 
   /// Parses tree in Ptr for base, index, offset addresses.
-  static BaseIndexOffset match(SDValue Ptr, const SelectionDAG &DAG);
+  static BaseIndexOffset match(LSBaseSDNode *N, const SelectionDAG &DAG);
 };
 
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_SELECTIONDAGADDRESSANALYSIS_H

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

for (unsigned i = 0; i < ByteWidth; i++) {
     // All loads must share the same chain
     SDValue LChain = L->getChain();
     if (!Chain)
       Chain = LChain;
     else if (Chain != LChain)
       return SDValue();
 
     // Loads must share the same base address
-    BaseIndexOffset Ptr = BaseIndexOffset::match(L->getBasePtr(), DAG);
+    BaseIndexOffset Ptr = BaseIndexOffset::match(L, DAG);
     int64_t ByteOffsetFromBase = 0;
     if (!Base)
       Base = Ptr;
     else if (!Base->equalBaseIndex(Ptr, DAG, ByteOffsetFromBase))
       return SDValue();
 
     // Calculate the offset of the current byte from the base address
     ByteOffsetFromBase += MemoryByteOffset(*P);
bool DAGCombiner::MergeStoresOfConstantsOrVecElts(
   AddToWorklist(NewChain.getNode());
   return true;
 }
 
 void DAGCombiner::getStoreMergeCandidates(
     StoreSDNode *St, SmallVectorImpl<MemOpLink> &StoreNodes) {
   // This holds the base pointer, index, and the offset in bytes from the base
   // pointer.
-  BaseIndexOffset BasePtr = BaseIndexOffset::match(St->getBasePtr(), DAG);
+  BaseIndexOffset BasePtr = BaseIndexOffset::match(St, DAG);
   EVT MemVT = St->getMemoryVT();
 
   SDValue Val = peekThroughBitcast(St->getValue());
   // We must have a base and an offset.
   if (!BasePtr.getBase().getNode())
     return;
 
   // Do not handle stores to undef base pointers.
   if (BasePtr.getBase().isUndef())
     return;
 
   bool IsConstantSrc = isa<ConstantSDNode>(Val) || isa<ConstantFPSDNode>(Val);
   bool IsExtractVecSrc = (Val.getOpcode() == ISD::EXTRACT_VECTOR_ELT ||
                           Val.getOpcode() == ISD::EXTRACT_SUBVECTOR);
   bool IsLoadSrc = isa<LoadSDNode>(Val);
   BaseIndexOffset LBasePtr;
   // Match on loadbaseptr if relevant.
   EVT LoadVT;
   if (IsLoadSrc) {
     auto *Ld = cast<LoadSDNode>(Val);
-    LBasePtr = BaseIndexOffset::match(Ld->getBasePtr(), DAG);
+    LBasePtr = BaseIndexOffset::match(Ld, DAG);
     LoadVT = Ld->getMemoryVT();
     // Load and store should be the same type.
     if (MemVT != LoadVT)
       return;
   }
   auto CandidateMatch = [&](StoreSDNode *Other, BaseIndexOffset &Ptr,
                             int64_t &Offset) -> bool {
     if (Other->isVolatile() || Other->isIndexed())
       return false;
     SDValue Val = peekThroughBitcast(Other->getValue());
     // Allow merging constants of different types as integers.
     bool NoTypeMatch = (MemVT.isInteger()) ? !MemVT.bitsEq(Other->getMemoryVT())
                                            : Other->getMemoryVT() != MemVT;
     if (IsLoadSrc) {
       if (NoTypeMatch)
         return false;
       // The Load's Base Ptr must also match
       if (LoadSDNode *OtherLd = dyn_cast<LoadSDNode>(Val)) {
-        auto LPtr = BaseIndexOffset::match(OtherLd->getBasePtr(), DAG);
+        auto LPtr = BaseIndexOffset::match(OtherLd, DAG);
         if (LoadVT != OtherLd->getMemoryVT())
           return false;
         if (!(LBasePtr.equalBaseIndex(LPtr, DAG)))
           return false;
       } else
         return false;
     }
     if (IsConstantSrc) {
       if (NoTypeMatch)
         return false;
       if (!(isa<ConstantSDNode>(Val) || isa<ConstantFPSDNode>(Val)))
         return false;
     }
     if (IsExtractVecSrc) {
       // Do not merge truncated stores here.
       if (Other->isTruncatingStore())
         return false;
       if (!MemVT.bitsEq(Val.getValueType()))
         return false;
       if (Val.getOpcode() != ISD::EXTRACT_VECTOR_ELT &&
           Val.getOpcode() != ISD::EXTRACT_SUBVECTOR)
         return false;
     }
-    Ptr = BaseIndexOffset::match(Other->getBasePtr(), DAG);
+    Ptr = BaseIndexOffset::match(Other, DAG);
     return (BasePtr.equalBaseIndex(Ptr, DAG, Offset));
   };
 
   // We looking for a root node which is an ancestor to all mergable
   // stores. We search up through a load, to our root and then down
   // through all children. For instance we will find Store{1,2,3} if
   // St is Store1, Store2. or Store3 where the root is not a load
   // which always true for nonvolatile ops. TODO: Expand
for (unsigned i = 0; i < NumConsecutiveStores; ++i) {
       // The memory operands must not be volatile.
       if (Ld->isVolatile() || Ld->isIndexed())
         break;
 
       // The stored memory type must be the same.
       if (Ld->getMemoryVT() != MemVT)
         break;
 
-      BaseIndexOffset LdPtr = BaseIndexOffset::match(Ld->getBasePtr(), DAG);
+      BaseIndexOffset LdPtr = BaseIndexOffset::match(Ld, DAG);
       // If this is not the first ptr that we check.
       int64_t LdOffset = 0;
       if (LdBasePtr.getBase().getNode()) {
         // The base ptr must be the same.
         if (!LdBasePtr.equalBaseIndex(LdPtr, DAG, LdOffset))
           break;
       } else {
         // Check that all other base pointers are the same as this one.
         LdBasePtr = LdPtr;
       }
 
if (LoadSDNode *Ld = dyn_cast<LoadSDNode>(Value)) {
     }
   }
 
   // Deal with elidable overlapping chained stores.
   if (StoreSDNode *ST1 = dyn_cast<StoreSDNode>(Chain))
     if (OptLevel != CodeGenOpt::None && ST->isUnindexed() &&
         ST1->isUnindexed() && !ST1->isVolatile() && ST1->hasOneUse() &&
         !ST1->getBasePtr().isUndef() && !ST->isVolatile()) {
-      BaseIndexOffset STBasePtr = BaseIndexOffset::match(ST->getBasePtr(), DAG);
-      BaseIndexOffset ST1BasePtr =
-          BaseIndexOffset::match(ST1->getBasePtr(), DAG);
+      BaseIndexOffset STBasePtr = BaseIndexOffset::match(ST, DAG);
+      BaseIndexOffset ST1BasePtr = BaseIndexOffset::match(ST1, DAG);
       unsigned STBytes = ST->getMemoryVT().getStoreSize();
       unsigned ST1Bytes = ST1->getMemoryVT().getStoreSize();
       int64_t PtrDiff;
       // If this is a store who's preceeding store to a subset of the same
       // memory and no one other node is chained to that store we can
       // effectively drop the store. Do not remove stores to undef as they may
       // be used as data sinks.
 
bool DAGCombiner::isAlias(LSBaseSDNode *Op0, LSBaseSDNode *Op1) const {
 
   if (Op1->isInvariant() && Op0->writeMem())
     return false;
 
   unsigned NumBytes0 = Op0->getMemoryVT().getStoreSize();
   unsigned NumBytes1 = Op1->getMemoryVT().getStoreSize();
 
   // Check for BaseIndexOffset matching.
-  BaseIndexOffset BasePtr0 = BaseIndexOffset::match(Op0->getBasePtr(), DAG);
-  BaseIndexOffset BasePtr1 = BaseIndexOffset::match(Op1->getBasePtr(), DAG);
+  BaseIndexOffset BasePtr0 = BaseIndexOffset::match(Op0, DAG);
+  BaseIndexOffset BasePtr1 = BaseIndexOffset::match(Op1, DAG);
   int64_t PtrDiff;
   if (BasePtr0.equalBaseIndex(BasePtr1, DAG, PtrDiff))
     return !((NumBytes0 <= PtrDiff) || (PtrDiff + NumBytes1 <= 0));
 
   // If both BasePtr0 and BasePtr1 are FrameIndexes, we will not be
   // able to calculate their relative offset if at least one arises
   // from an alloca. However, these allocas cannot overlap and we
   // can infer there is no alias.

 // to go from a partially-merged state to the desired final
 // fully-merged state.
 bool DAGCombiner::findBetterNeighborChains(StoreSDNode *St) {
   if (OptLevel == CodeGenOpt::None)
     return false;
 
   // This holds the base pointer, index, and the offset in bytes from the base
   // pointer.
-  BaseIndexOffset BasePtr = BaseIndexOffset::match(St->getBasePtr(), DAG);
+  BaseIndexOffset BasePtr = BaseIndexOffset::match(St, DAG);
 
   // We must have a base and an offset.
   if (!BasePtr.getBase().getNode())
     return false;
 
   // Do not handle stores to undef base pointers.
   if (BasePtr.getBase().isUndef())
     return false;
while (Index) {
     // If the chain has more than one use, then we can't reorder the mem ops.
     if (Index != St && !SDValue(Index, 0)->hasOneUse())
       break;
 
     if (Index->isVolatile() || Index->isIndexed())
       break;
 
     // Find the base pointer and offset for this memory node.
-    BaseIndexOffset Ptr = BaseIndexOffset::match(Index->getBasePtr(), DAG);
+    BaseIndexOffset Ptr = BaseIndexOffset::match(Index, DAG);
 
     // Check that the base pointer is the same as the original one.
     if (!BasePtr.equalBaseIndex(Ptr, DAG))
       break;
 
     // Walk up the chain to find the next store node, ignoring any
     // intermediate loads. Any other kind of node will halt the loop.
     SDNode *NextInChain = Index->getChain().getNode();

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

bool SelectionDAG::areNonVolatileConsecutiveLoads(LoadSDNode *LD,
   if (LD->isIndexed() || Base->isIndexed())
     return false;
   if (LD->getChain() != Base->getChain())
     return false;
   EVT VT = LD->getValueType(0);
   if (VT.getSizeInBits() / 8 != Bytes)
     return false;
 
-  SDValue Loc = LD->getOperand(1);
-  SDValue BaseLoc = Base->getOperand(1);
-
-  auto BaseLocDecomp = BaseIndexOffset::match(BaseLoc, *this);
-  auto LocDecomp = BaseIndexOffset::match(Loc, *this);
+  auto BaseLocDecomp = BaseIndexOffset::match(Base, *this);
+  auto LocDecomp = BaseIndexOffset::match(LD, *this);
 
   int64_t Offset = 0;
   if (BaseLocDecomp.equalBaseIndex(LocDecomp, *this, Offset))
     return (Dist * Bytes == Offset);
   return false;
 }
 
 /// InferPtrAlignment - Infer alignment of a load / store address. Return 0 if

llvm/lib/CodeGen/SelectionDAG/SelectionDAGAddressAnalysis.cpp

if (auto *A = dyn_cast<FrameIndexSDNode>(Base))
                  MFI.getObjectOffset(A->getIndex());
           return true;
         }
   }
   return false;
 }
 
 /// Parses tree in Ptr for base, index, offset addresses.
-BaseIndexOffset BaseIndexOffset::match(SDValue Ptr, const SelectionDAG &DAG) {
+BaseIndexOffset BaseIndexOffset::match(LSBaseSDNode *N,
+                                       const SelectionDAG &DAG) {
+  SDValue Ptr = N->getBasePtr();
+
   // (((B + I*M) + c)) + c ...
   SDValue Base = DAG.getTargetLoweringInfo().unwrapAddress(Ptr);
   SDValue Index = SDValue();
   int64_t Offset = 0;
   bool IsIndexSignExt = false;
 
+  // pre-inc/pre-dec ops are components of EA.
+  if (N->getAddressingMode() == ISD::PRE_INC) {
+    if (auto *C = dyn_cast<ConstantSDNode>(N->getOffset()))
+      Offset += C->getSExtValue();
+    else // If unknown, give up now.
+      return BaseIndexOffset(Base, Index, Offset, IsIndexSignExt);
+  } else if (N->getAddressingMode() == ISD::PRE_DEC) {
+    if (auto *C = dyn_cast<ConstantSDNode>(N->getOffset()))
+      Offset -= C->getSExtValue();
+    else // If unknown, give up now.
+      return BaseIndexOffset(Base, Index, Offset, IsIndexSignExt);
+  }
+
   // Consume constant adds & ors with appropriate masking.
   while (Base->getOpcode() == ISD::ADD || Base->getOpcode() == ISD::OR) {
     if (auto *C = dyn_cast<ConstantSDNode>(Base->getOperand(1))) {
       // Only consider ORs which act as adds.
       if (Base->getOpcode() == ISD::OR &&
           !DAG.MaskedValueIsZero(Base->getOperand(0), C->getAPIntValue()))
         break;
       Offset += C->getSExtValue();