[LSV] Look through selects for consecutive addresses

[rtereshin]

Summary

In some cases LSV sees (load/store _ (select _ <pointer expression> <pointer expression>))
patterns in input IR, often due to sinking and other forms of CFG simplification, sometimes
interspersed with bitcasts and all-constant-indices GEPs. With this patch`areConsecutivePointers`
method would attempt to handle select instructions. This leads to an increased number of
successful vectorizations.

Technically, select instructions could appear in index arithmetic as well, however, we don't see
those in our test suites / benchmarks. Also, there is a lot more freedom in IR shapes computing
integral indices in general than in what's common in pointer computations, and it appears that
it's quite unreliable to do anything short of making select instructions first class citizens of
Scalar Evolution, which for the purposes of this patch is most definitely an overkill.

[rampitec]

LGTM

[arsenm]

lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

147

Why is this an enum?

459–460

Braces

[arsenm]

lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

322

Const method?

695–700

I'm a bit confused by this ChainID concept. Why is the select condition used for this and not the select itself?

[rtereshin]

lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

147

Sure, will change that to static const or static constexpr. Which one would you prefer?

322

Maybe, I never checked. Most likely not, it calls methods on ScalarEvolution &SE, which is part of the object, and those methods change the SCEVs' cache that is owned by SE, so most likely they aren't const methods themselves.

If it is a const method though, will do.

459–460

Where exactly would you like braces? First level if, second level if, or both?

I believe no curly braces around 1 statement bodies of if/then/else/for/while/do-while statements is the standard LLVM code style.

695–700

I'm a bit confused by this ChainID concept.

ChainID is an arbitrary token that is allowed to be different only for the accesses that are guaranteed to be non-consecutive. It's used (and it's predecessor was used) for grouping instructions together and reducing the number of instructions the main search operates on at a time, i.e. this is to reduce compile time and nothing else as the main search has (and had) O(n^2) time complexity.

The exact way of determining accesses that are guaranteed to be non-consecutive does not affect the rest of the pass as soon as it doesn't return different ChainIDs for consecutive accesses. I played with a few different ones and I expect it to change in the future. All one needs to do to change it is to change the type definition and the getChainID function implementation. In particular, returning the same ChainID for every Ptr works too, but increases compile time.

Of course, "guaranteed to be non-consecutive" above means "guaranteed to be classified as non-consecutive by isConsecutiveAccess method".

Why is the select condition used for this and not the select itself?

The select's themselves are distinct instructions even if they share the same condition and evaluate to consecutive pointers for true and false values of the condition. Therefore using the select's themselves for grouping instructions would put consecutive accesses into different lists and they wouldn't be even checked for being consecutive, and won't be vectorized. Which defies the purpose of this patch.

[rtereshin]

lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

147

Going with static const for now.

322

Could be done, good catch, doing that.

459–460

Going with braces around second level if for now.

[rtereshin]

1. Turned methods touched to const methods
2. Added curly braces
3. Turned single-value enum into static const

[rtereshin]

@arsenm

Hi Matt,

Do you think this is good to go now?

Thanks,
Roman

[arsenm]

lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

459–460

I would put both, as both are surrounding multi-line things

695–700

This should be in a comment somewhere

[rtereshin]

1. Added comments explaining partitioning of memory access instructions into groups of candidates for better performance
2. Added curly braces

lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

459–460

Done.

695–700

Done.

[volkan]

Hi Roman,

I added a few comments. Could you take a look?

lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

147

Replace RecursionLimit with MaxDepth maybe?

323

Replace MaxRecurse with Depth maybe?

360

Missing --MaxRecurse?

368

You can check if isa<SelectInst>(A) && isa<SelectInst>(B) here.

454

You can replace Value * with SelectInst & and remove the dyn_casts below.

[rtereshin]

lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

323

ok

360

no, lookThroughComplexAddresses does not make recursive calls directly, lookThroughSelects does later.

368

I can, but it gets messier.

454

I can, but it gets messier on the caller's side.

[rtereshin]

1. Renamed MaxRecurse to Depth as requested and updated the logic around it as MaxRecurse was going from the positive limit down to 0 which makes no sense for a parameter called Depth, so Depth is going from 0 to MaxDepth now.
2. Renamed RecursionLimit to MaxDepth to make the connection with Depth parameter clearer (which was MaxRecurse previously).

[volkan]

LGTM

[rtereshin]

In D49428#1175824, @volkan wrote:

LGTM

Thanks!

@arsenm Hopefully it's fine to move forward with this one at this point.

Roman

lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

 STATISTIC(NumVectorInstructions, "Number of vector accesses generated");
 STATISTIC(NumScalarsVectorized, "Number of scalar accesses vectorized");
 
 // FIXME: Assuming stack alignment of 4 is always good enough
 static const unsigned StackAdjustedAlignment = 4;
 
 namespace {
 
+/// ChainID is an arbitrary token that is allowed to be different only for the
+/// accesses that are guaranteed to be considered non-consecutive by
+/// Vectorizer::isConsecutiveAccess. It's used for grouping instructions
+/// together and reducing the number of instructions the main search operates on
+/// at a time, i.e. this is to reduce compile time and nothing else as the main
+/// search has O(n^2) time complexity. The underlying type of ChainID should not
+/// be relied upon.
+using ChainID = const Value *;
 using InstrList = SmallVector<Instruction *, 8>;
-using InstrListMap = MapVector<Value *, InstrList>;
+using InstrListMap = MapVector<ChainID, InstrList>;
 
 class Vectorizer {
   Function &F;
   AliasAnalysis &AA;
   DominatorTree &DT;
   ScalarEvolution &SE;
   TargetTransformInfo &TTI;
   const DataLayout &DL;
private:
   unsigned getAlignment(StoreInst *SI) const {
     unsigned Align = SI->getAlignment();
     if (Align != 0)
       return Align;
 
     return DL.getABITypeAlignment(SI->getValueOperand()->getType());
   }
 
+  static const unsigned RecursionLimit = 3;
+
   bool isConsecutiveAccess(Value *A, Value *B);
-  bool areConsecutivePointers(Value *PtrA, Value *PtrB, APInt Size);
-  bool lookThroughComplexAddresses(Value *PtrA, Value *PtrB, APInt PtrDelta);
+  bool areConsecutivePointers(Value *PtrA, Value *PtrB, APInt PtrDelta,
+                              unsigned MaxRecurse) const;
+  bool lookThroughComplexAddresses(Value *PtrA, Value *PtrB, APInt PtrDelta,
+                                   unsigned MaxRecurse) const;
+  bool lookThroughSelects(Value *PtrA, Value *PtrB, APInt PtrDelta,
+                          unsigned MaxRecurse) const;
 
   /// After vectorization, reorder the instructions that I depends on
   /// (the instructions defining its operands), to ensure they dominate I.
   void reorder(Instruction *I);
 
   /// Returns the first and the last instructions in Chain.
   std::pair<BasicBlock::iterator, BasicBlock::iterator>
   getBoundaryInstrs(ArrayRef<Instruction *> Chain);
if (PtrA == PtrB ||
       DL.getTypeStoreSize(PtrATy) != DL.getTypeStoreSize(PtrBTy) ||
       DL.getTypeStoreSize(PtrATy->getScalarType()) !=
           DL.getTypeStoreSize(PtrBTy->getScalarType()))
     return false;
 
   unsigned PtrBitWidth = DL.getPointerSizeInBits(ASA);
   APInt Size(PtrBitWidth, DL.getTypeStoreSize(PtrATy));
 
-  return areConsecutivePointers(PtrA, PtrB, Size);
+  return areConsecutivePointers(PtrA, PtrB, Size, RecursionLimit);
 }
 
-bool Vectorizer::areConsecutivePointers(Value *PtrA, Value *PtrB, APInt Size) {
+bool Vectorizer::areConsecutivePointers(Value *PtrA, Value *PtrB,
+                                        APInt PtrDelta,
+                                        unsigned MaxRecurse) const {
   unsigned PtrBitWidth = DL.getPointerTypeSizeInBits(PtrA->getType());
   APInt OffsetA(PtrBitWidth, 0);
   APInt OffsetB(PtrBitWidth, 0);
   PtrA = PtrA->stripAndAccumulateInBoundsConstantOffsets(DL, OffsetA);
   PtrB = PtrB->stripAndAccumulateInBoundsConstantOffsets(DL, OffsetB);
 
   APInt OffsetDelta = OffsetB - OffsetA;
 
   // Check if they are based on the same pointer. That makes the offsets
   // sufficient.
   if (PtrA == PtrB)
-    return OffsetDelta == Size;
+    return OffsetDelta == PtrDelta;
 
   // Compute the necessary base pointer delta to have the necessary final delta
-  // equal to the size.
-  APInt BaseDelta = Size - OffsetDelta;
+  // equal to the pointer delta requested.
+  APInt BaseDelta = PtrDelta - OffsetDelta;
 
   // Compute the distance with SCEV between the base pointers.
   const SCEV *PtrSCEVA = SE.getSCEV(PtrA);
   const SCEV *PtrSCEVB = SE.getSCEV(PtrB);
   const SCEV *C = SE.getConstant(BaseDelta);
   const SCEV *X = SE.getAddExpr(PtrSCEVA, C);
   if (X == PtrSCEVB)
     return true;
 
   // The above check will not catch the cases where one of the pointers is
   // factorized but the other one is not, such as (C + (S * (A + B))) vs
   // (AS + BS). Get the minus scev. That will allow re-combining the expresions
   // and getting the simplified difference.
   const SCEV *Dist = SE.getMinusSCEV(PtrSCEVB, PtrSCEVA);
   if (C == Dist)
     return true;
 
   // Sometimes even this doesn't work, because SCEV can't always see through
   // patterns that look like (gep (ext (add (shl X, C1), C2))). Try checking
   // things the hard way.
-  return lookThroughComplexAddresses(PtrA, PtrB, BaseDelta);
+  return lookThroughComplexAddresses(PtrA, PtrB, BaseDelta, MaxRecurse);
 }
 
 bool Vectorizer::lookThroughComplexAddresses(Value *PtrA, Value *PtrB,
-                                             APInt PtrDelta) {
+                                             APInt PtrDelta,
+                                             unsigned MaxRecurse) const {
   auto *GEPA = dyn_cast<GetElementPtrInst>(PtrA);
   auto *GEPB = dyn_cast<GetElementPtrInst>(PtrB);
   if (!GEPA || !GEPB)
-    return false;
+    return lookThroughSelects(PtrA, PtrB, PtrDelta, MaxRecurse);
 
   // Look through GEPs after checking they're the same except for the last
   // index.
   if (GEPA->getNumOperands() != GEPB->getNumOperands() ||
       GEPA->getPointerOperand() != GEPB->getPointerOperand())
     return false;
   gep_type_iterator GTIA = gep_type_begin(GEPA);
   gep_type_iterator GTIB = gep_type_begin(GEPB);
bool Vectorizer::lookThroughComplexAddresses(Value *PtrA, Value *PtrB,
 
   const SCEV *OffsetSCEVA = SE.getSCEV(ValA);
   const SCEV *OffsetSCEVB = SE.getSCEV(OpB);
   const SCEV *C = SE.getConstant(IdxDiff.trunc(BitWidth));
   const SCEV *X = SE.getAddExpr(OffsetSCEVA, C);
   return X == OffsetSCEVB;
 }
 
+bool Vectorizer::lookThroughSelects(Value *PtrA, Value *PtrB, APInt PtrDelta,
+                                    unsigned MaxRecurse) const {
+  if (!MaxRecurse--)
+    return false;
+
+  if (auto *SelectA = dyn_cast<SelectInst>(PtrA)) {
+    if (auto *SelectB = dyn_cast<SelectInst>(PtrB)) {
+      return SelectA->getCondition() == SelectB->getCondition() &&
+             areConsecutivePointers(SelectA->getTrueValue(),
+                                    SelectB->getTrueValue(), PtrDelta,
+                                    MaxRecurse) &&
+             areConsecutivePointers(SelectA->getFalseValue(),
+                                    SelectB->getFalseValue(), PtrDelta,
+                                    MaxRecurse);
+    }
+  }
+  return false;
+}
+
 void Vectorizer::reorder(Instruction *I) {
   OrderedBasicBlock OBB(I->getParent());
   SmallPtrSet<Instruction *, 16> InstructionsToMove;
   SmallVector<Instruction *, 16> Worklist;
 
   Worklist.push_back(I);
   while (!Worklist.empty()) {
     Instruction *IW = Worklist.pop_back_val();
Vectorizer::getVectorizablePrefix(ArrayRef<Instruction *> Chain) {
   unsigned ChainIdx = 0;
   for (unsigned ChainLen = Chain.size(); ChainIdx < ChainLen; ++ChainIdx) {
     if (!VectorizableChainInstrs.count(Chain[ChainIdx]))
       break;
   }
   return Chain.slice(0, ChainIdx);
 }
 
+static ChainID getChainID(const Value *Ptr, const DataLayout &DL) {
+  const Value *ObjPtr = GetUnderlyingObject(Ptr, DL);
+  if (const auto *Sel = dyn_cast<SelectInst>(ObjPtr)) {
+    // The select's themselves are distinct instructions even if they share the
+    // same condition and evaluate to consecutive pointers for true and false
+    // values of the condition. Therefore using the select's themselves for
+    // grouping instructions would put consecutive accesses into different lists
+    // and they won't be even checked for being consecutive, and won't be
+    // vectorized.
+    return Sel->getCondition();
+  }
+  return ObjPtr;
+}
+
 std::pair<InstrListMap, InstrListMap>
 Vectorizer::collectInstructions(BasicBlock *BB) {
   InstrListMap LoadRefs;
   InstrListMap StoreRefs;
 
   for (Instruction &I : *BB) {
     if (!I.mayReadOrWriteMemory())
       continue;
if (LoadInst *LI = dyn_cast<LoadInst>(&I)) {
       // Make sure all the users of a vector are constant-index extracts.
       if (isa<VectorType>(Ty) && !llvm::all_of(LI->users(), [](const User *U) {
             const ExtractElementInst *EEI = dyn_cast<ExtractElementInst>(U);
             return EEI && isa<ConstantInt>(EEI->getOperand(1));
           }))
         continue;
 
       // Save the load locations.
-      Value *ObjPtr = GetUnderlyingObject(Ptr, DL);
-      LoadRefs[ObjPtr].push_back(LI);
+      const ChainID ID = getChainID(Ptr, DL);
+      LoadRefs[ID].push_back(LI);
     } else if (StoreInst *SI = dyn_cast<StoreInst>(&I)) {
       if (!SI->isSimple())
         continue;
 
       // Skip if it's not legal.
       if (!TTI.isLegalToVectorizeStore(SI))
         continue;
 
if (LoadInst *LI = dyn_cast<LoadInst>(&I)) {
 
       if (isa<VectorType>(Ty) && !llvm::all_of(SI->users(), [](const User *U) {
             const ExtractElementInst *EEI = dyn_cast<ExtractElementInst>(U);
             return EEI && isa<ConstantInt>(EEI->getOperand(1));
           }))
         continue;
 
       // Save store location.
-      Value *ObjPtr = GetUnderlyingObject(Ptr, DL);
-      StoreRefs[ObjPtr].push_back(SI);
+      const ChainID ID = getChainID(Ptr, DL);
+      StoreRefs[ID].push_back(SI);
     }
   }
 
   return {LoadRefs, StoreRefs};
 }
 
 bool Vectorizer::vectorizeChains(InstrListMap &Map) {
   bool Changed = false;
 
-  for (const std::pair<Value *, InstrList> &Chain : Map) {
+  for (const std::pair<ChainID, InstrList> &Chain : Map) {
     unsigned Size = Chain.second.size();
     if (Size < 2)
       continue;
 
     LLVM_DEBUG(dbgs() << "LSV: Analyzing a chain of length " << Size << ".\n");
 
     // Process the stores in chunks of 64.
     for (unsigned CI = 0, CE = Size; CI < CE; CI += 64) {

test/Transforms/LoadStoreVectorizer/AMDGPU/selects.ll

+; RUN: opt -mtriple=amdgcn-amd-amdhsa -load-store-vectorizer -dce -S -o - %s | FileCheck %s
+
+target datalayout = "e-p:32:32-p1:64:64-p2:64:64-p3:32:32-p4:64:64-p5:32:32-p24:64:64-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64"
+
+define void @base_case(i1 %cnd, i32 addrspace(1)* %a, i32 addrspace(1)* %b, <3 x i32> addrspace(1)* %out) {
+; CHECK-LABEL: @base_case
+; CHECK: load <3 x i32>
+entry:
+  %gep1 = getelementptr inbounds i32, i32 addrspace(1)* %a, i64 1
+  %gep2 = getelementptr inbounds i32, i32 addrspace(1)* %a, i64 2
+  %gep4 = getelementptr inbounds i32, i32 addrspace(1)* %b, i64 1
+  %gep5 = getelementptr inbounds i32, i32 addrspace(1)* %b, i64 2
+  %selected = select i1 %cnd, i32 addrspace(1)* %a, i32 addrspace(1)* %b
+  %selected14 = select i1 %cnd, i32 addrspace(1)* %gep1, i32 addrspace(1)* %gep4
+  %selected25 = select i1 %cnd, i32 addrspace(1)* %gep2, i32 addrspace(1)* %gep5
+  %val0 = load i32, i32 addrspace(1)* %selected, align 4
+  %val1 = load i32, i32 addrspace(1)* %selected14, align 4
+  %val2 = load i32, i32 addrspace(1)* %selected25, align 4
+  %t0 = insertelement <3 x i32> undef, i32 %val0, i32 0
+  %t1 = insertelement <3 x i32> %t0, i32 %val1, i32 1
+  %t2 = insertelement <3 x i32> %t1, i32 %val2, i32 2
+  store <3 x i32> %t2, <3 x i32> addrspace(1)* %out
+  ret void
+}
+
+define void @scev_targeting_complex_case(i1 %cnd, i32 addrspace(1)* %a, i32 addrspace(1)* %b, i32 %base, <2 x i32> addrspace(1)* %out) {
+; CHECK-LABEL: @scev_targeting_complex_case
+; CHECK: load <2 x i32>
+entry:
+  %base.x4 = shl i32 %base, 2
+  %base.x4.p1 = add i32 %base.x4, 1
+  %base.x4.p2 = add i32 %base.x4, 2
+  %base.x4.p3 = add i32 %base.x4, 3
+  %zext.x4 = zext i32 %base.x4 to i64
+  %zext.x4.p1 = zext i32 %base.x4.p1 to i64
+  %zext.x4.p2 = zext i32 %base.x4.p2 to i64
+  %zext.x4.p3 = zext i32 %base.x4.p3 to i64
+  %base.x16 = mul i64 %zext.x4, 4
+  %base.x16.p4 = shl i64 %zext.x4.p1, 2
+  %base.x16.p8 = shl i64 %zext.x4.p2, 2
+  %base.x16.p12 = mul i64 %zext.x4.p3, 4
+  %a.pi8 = bitcast i32 addrspace(1)* %a to i8 addrspace(1)*
+  %b.pi8 = bitcast i32 addrspace(1)* %b to i8 addrspace(1)*
+  %gep.a.base.x16 = getelementptr inbounds i8, i8 addrspace(1)* %a.pi8, i64 %base.x16
+  %gep.b.base.x16.p4 = getelementptr inbounds i8, i8 addrspace(1)* %b.pi8, i64 %base.x16.p4
+  %gep.a.base.x16.p8 = getelementptr inbounds i8, i8 addrspace(1)* %a.pi8, i64 %base.x16.p8
+  %gep.b.base.x16.p12 = getelementptr inbounds i8, i8 addrspace(1)* %b.pi8, i64 %base.x16.p12
+  %a.base.x16 = bitcast i8 addrspace(1)* %gep.a.base.x16 to i32 addrspace(1)*
+  %b.base.x16.p4 = bitcast i8 addrspace(1)* %gep.b.base.x16.p4 to i32 addrspace(1)*
+  %selected.base.x16.p0.or.4 = select i1 %cnd, i32 addrspace(1)* %a.base.x16, i32 addrspace(1)* %b.base.x16.p4
+  %gep.selected.base.x16.p8.or.12 = select i1 %cnd, i8 addrspace(1)* %gep.a.base.x16.p8, i8 addrspace(1)* %gep.b.base.x16.p12
+  %selected.base.x16.p8.or.12 = bitcast i8 addrspace(1)* %gep.selected.base.x16.p8.or.12 to i32 addrspace(1)*
+  %selected.base.x16.p40.or.44 = getelementptr inbounds i32, i32 addrspace(1)* %selected.base.x16.p0.or.4, i64 10
+  %selected.base.x16.p44.or.48 = getelementptr inbounds i32, i32 addrspace(1)* %selected.base.x16.p8.or.12, i64 9
+  %val0 = load i32, i32 addrspace(1)* %selected.base.x16.p40.or.44, align 4
+  %val1 = load i32, i32 addrspace(1)* %selected.base.x16.p44.or.48, align 4
+  %t0 = insertelement <2 x i32> undef, i32 %val0, i32 0
+  %t1 = insertelement <2 x i32> %t0, i32 %val1, i32 1
+  store <2 x i32> %t1, <2 x i32> addrspace(1)* %out
+  ret void
+}
+
+define void @nested_selects(i1 %cnd0, i1 %cnd1, i32 addrspace(1)* %a, i32 addrspace(1)* %b, i32 %base, <2 x i32> addrspace(1)* %out) {
+; CHECK-LABEL: @nested_selects
+; CHECK: load <2 x i32>
+entry:
+  %base.p1 = add nsw i32 %base, 1
+  %base.p2 = add i32 %base, 2
+  %base.p3 = add nsw i32 %base, 3
+  %base.x4 = mul i32 %base, 4
+  %base.x4.p5 = add i32 %base.x4, 5
+  %base.x4.p6 = add i32 %base.x4, 6
+  %sext = sext i32 %base to i64
+  %sext.p1 = sext i32 %base.p1 to i64
+  %sext.p2 = sext i32 %base.p2 to i64
+  %sext.p3 = sext i32 %base.p3 to i64
+  %sext.x4.p5 = sext i32 %base.x4.p5 to i64
+  %sext.x4.p6 = sext i32 %base.x4.p6 to i64
+  %gep.a.base = getelementptr inbounds i32, i32 addrspace(1)* %a, i64 %sext
+  %gep.a.base.p1 = getelementptr inbounds i32, i32 addrspace(1)* %a, i64 %sext.p1
+  %gep.a.base.p2 = getelementptr inbounds i32, i32 addrspace(1)* %a, i64 %sext.p2
+  %gep.a.base.p3 = getelementptr inbounds i32, i32 addrspace(1)* %a, i64 %sext.p3
+  %gep.b.base.x4.p5 = getelementptr inbounds i32, i32 addrspace(1)* %a, i64 %sext.x4.p5
+  %gep.b.base.x4.p6 = getelementptr inbounds i32, i32 addrspace(1)* %a, i64 %sext.x4.p6
+  %selected.1.L = select i1 %cnd1, i32 addrspace(1)* %gep.a.base.p2, i32 addrspace(1)* %gep.b.base.x4.p5
+  %selected.1.R = select i1 %cnd1, i32 addrspace(1)* %gep.a.base.p3, i32 addrspace(1)* %gep.b.base.x4.p6
+  %selected.0.L = select i1 %cnd0, i32 addrspace(1)* %gep.a.base, i32 addrspace(1)* %selected.1.L
+  %selected.0.R = select i1 %cnd0, i32 addrspace(1)* %gep.a.base.p1, i32 addrspace(1)* %selected.1.R
+  %val0 = load i32, i32 addrspace(1)* %selected.0.L, align 4
+  %val1 = load i32, i32 addrspace(1)* %selected.0.R, align 4
+  %t0 = insertelement <2 x i32> undef, i32 %val0, i32 0
+  %t1 = insertelement <2 x i32> %t0, i32 %val1, i32 1
+  store <2 x i32> %t1, <2 x i32> addrspace(1)* %out
+  ret void
+}