[VectorCombine] Use constant range info for index scalarization legality.

[fhahn]

Summary

We can only scalarize memory accesses if we know the index is valid.

This patch adjusts canScalarizeAcceess to fall back to
computeConstantRange to check if the index is known to be valid.

[fhahn]

Also get and pass in the AC with the legacy PM.

[RKSimon]

What about an index that has been clamped with a urem() or and() ?

[fhahn]

rebase after adding positive & negative tests with urem & and in 2f69b78a578d

[fhahn]

rebase & ping :)

[nlopes]

This patch looks good if we want to go with the poison-on-OOB semantics.
I think it's important to settle that first, please see my bug report: https://bugs.llvm.org/show_bug.cgi?id=50382
The backend does this very same optimization but without paying attention to whether the index is in-bounds or not. Which effectively makes extractvalue similar to a load instruction, including all the restrictions around movement and so on. Is that what we want?
We need to either fix the backend and move on with this patch (poison on OOB), or fix langref and abandon this patch (UB on OOB). Since I don't really know which frontend produces extractelement and from which code, I don't have an opinion either way, although usually UB is not a good idea.

[fhahn]

In D102476#2777021, @nlopes wrote:

This patch looks good if we want to go with the poison-on-OOB semantics.
I think it's important to settle that first, please see my bug report: https://bugs.llvm.org/show_bug.cgi?id=50382

I think the conclusion of the discussion at the patch to far seems to be to keep it poison. So I think we could move ahead with this patch?

[nlopes]

In D102476#2779213, @fhahn wrote:

In D102476#2777021, @nlopes wrote:

This patch looks good if we want to go with the poison-on-OOB semantics.
I think it's important to settle that first, please see my bug report: https://bugs.llvm.org/show_bug.cgi?id=50382

I think the conclusion of the discussion at the patch to far seems to be to keep it poison. So I think we could move ahead with this patch?

SGTM!

[efriedma]

llvm/test/Transforms/VectorCombine/load-insert-store.ll

190

Alive doesn't like this: https://web.ist.utl.pt/nuno.lopes/alive2/index.php?hash=50de39e1c7237e13&test=Transforms%2FVectorCombine%2Fload-insert-store.ll

[efriedma]

llvm/test/Transforms/VectorCombine/load-insert-store.ll

190

Filed as https://bugs.llvm.org/show_bug.cgi?id=50949

[fhahn]

llvm/test/Transforms/VectorCombine/load-insert-store.ll

190

Apologies for the delay! I put up a patch limiting the transform to cases where the index is known non-poison: D107364

[fhahn]

llvm/test/Transforms/VectorCombine/load-insert-store.ll

190

Ok, looks like things look good now on the latest run https://web.ist.utl.pt/nuno.lopes/alive2/index.php?hash=295f040c6c7caf3e

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

 //===------- VectorCombine.cpp - Optimize partial vector operations -------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This pass optimizes scalar/vector interactions using target cost models. The
 // transforms implemented here may not fit in traditional loop-based or SLP
 // vectorization passes.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Vectorize/VectorCombine.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/BasicAliasAnalysis.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/Loads.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/Analysis/VectorUtils.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
static cl::opt<unsigned> MaxInstrsToScan(
     cl::desc("Max number of instructions to scan for vector combining."));
 
 static const unsigned InvalidIndex = std::numeric_limits<unsigned>::max();
 
 namespace {
 class VectorCombine {
 public:
   VectorCombine(Function &F, const TargetTransformInfo &TTI,
-                const DominatorTree &DT, AAResults &AA)
-      : F(F), Builder(F.getContext()), TTI(TTI), DT(DT), AA(AA) {}
+                const DominatorTree &DT, AAResults &AA,
+                AssumptionCache *AC = nullptr)
+      : F(F), Builder(F.getContext()), TTI(TTI), DT(DT), AA(AA), AC(AC) {}
 
   bool run();
 
 private:
   Function &F;
   IRBuilder<> Builder;
   const TargetTransformInfo &TTI;
   const DominatorTree &DT;
   AAResults &AA;
+  AssumptionCache *AC;
 
   bool vectorizeLoadInsert(Instruction &I);
   ExtractElementInst *getShuffleExtract(ExtractElementInst *Ext0,
                                         ExtractElementInst *Ext1,
                                         unsigned PreferredExtractIndex) const;
   bool isExtractExtractCheap(ExtractElementInst *Ext0, ExtractElementInst *Ext1,
                              unsigned Opcode,
                              ExtractElementInst *&ConvertToShuffle,
static bool isMemModifiedBetween(BasicBlock::iterator Begin,
   return std::any_of(Begin, End, [&](const Instruction &Instr) {
     return isModSet(AA.getModRefInfo(&Instr, Loc)) ||
            ++NumScanned > MaxInstrsToScan;
   });
 }
 
 /// Check if it is legal to scalarize a memory access to \p VecTy at index \p
 /// Idx. \p Idx must access a valid vector element.
-static bool canScalarizeAccess(FixedVectorType *VecTy, ConstantInt *Idx) {
-  return Idx->getValue().ult(VecTy->getNumElements());
+static bool canScalarizeAccess(FixedVectorType *VecTy, Value *Idx,
+                               Instruction *CtxI, AssumptionCache *AC) {
+  if (auto *C = dyn_cast<ConstantInt>(Idx))
+    return C->getValue().ult(VecTy->getNumElements());
+
+  APInt Zero(Idx->getType()->getScalarSizeInBits(), 0);
+  APInt MaxElts(Idx->getType()->getScalarSizeInBits(), VecTy->getNumElements());
+  ConstantRange ValidIndices(Zero, MaxElts);
+  ConstantRange IdxRange = computeConstantRange(Idx, true, AC, CtxI, 0);
+  return ValidIndices.contains(IdxRange);
 }
 
 // Combine patterns like:
 //   %0 = load <4 x i32>, <4 x i32>* %a
 //   %1 = insertelement <4 x i32> %0, i32 %b, i32 1
 //   store <4 x i32> %1, <4 x i32>* %a
 // to:
 //   %0 = bitcast <4 x i32>* %a to i32*
 //   %1 = getelementptr inbounds i32, i32* %0, i64 0, i64 1
 //   store i32 %b, i32* %1
 bool VectorCombine::foldSingleElementStore(Instruction &I) {
   StoreInst *SI = dyn_cast<StoreInst>(&I);
   if (!SI || !SI->isSimple() ||
       !isa<FixedVectorType>(SI->getValueOperand()->getType()))
     return false;
 
   // TODO: Combine more complicated patterns (multiple insert) by referencing
   // TargetTransformInfo.
   Instruction *Source;
   Value *NewElement;
-  ConstantInt *Idx;
+  Value *Idx;
   if (!match(SI->getValueOperand(),
              m_InsertElt(m_Instruction(Source), m_Value(NewElement),
-                         m_ConstantInt(Idx))))
+                         m_Value(Idx))))
     return false;
 
   if (auto *Load = dyn_cast<LoadInst>(Source)) {
     auto VecTy = cast<FixedVectorType>(SI->getValueOperand()->getType());
     const DataLayout &DL = I.getModule()->getDataLayout();
     Value *SrcAddr = Load->getPointerOperand()->stripPointerCasts();
     // Don't optimize for atomic/volatile load or store. Ensure memory is not
     // modified between, vector type matches store size, and index is inbounds.
     if (!Load->isSimple() || Load->getParent() != SI->getParent() ||
         !DL.typeSizeEqualsStoreSize(Load->getType()) ||
-        !canScalarizeAccess(VecTy, Idx) ||
+        !canScalarizeAccess(VecTy, Idx, Load, AC) ||
         SrcAddr != SI->getPointerOperand()->stripPointerCasts() ||
         isMemModifiedBetween(Load->getIterator(), SI->getIterator(),
                              MemoryLocation::get(SI), AA))
       return false;
 
     Value *GEP = GetElementPtrInst::CreateInBounds(
         SI->getPointerOperand(), {ConstantInt::get(Idx->getType(), 0), Idx});
     Builder.Insert(GEP);
     StoreInst *NSI = Builder.CreateStore(NewElement, GEP);
     NSI->copyMetadata(*SI);
     if (SI->getAlign() < NSI->getAlign())
       NSI->setAlignment(SI->getAlign());
     replaceValue(I, *NSI);
     // Need erasing the store manually.
     I.eraseFromParent();
     return true;
   }
 
   return false;
 }
 
 /// Try to scalarize vector loads feeding extractelement instructions.
 bool VectorCombine::scalarizeLoadExtract(Instruction &I) {
   Value *Ptr;
-  ConstantInt *Idx;
-  if (!match(&I, m_ExtractElt(m_Load(m_Value(Ptr)), m_ConstantInt(Idx))))
+  Value *Idx;
+  if (!match(&I, m_ExtractElt(m_Load(m_Value(Ptr)), m_Value(Idx))))
     return false;
 
   auto *LI = cast<LoadInst>(I.getOperand(0));
   const DataLayout &DL = I.getModule()->getDataLayout();
   if (LI->isVolatile() || !DL.typeSizeEqualsStoreSize(LI->getType()))
     return false;
 
   auto *FixedVT = dyn_cast<FixedVectorType>(LI->getType());
   if (!FixedVT)
     return false;
 
-  if (!canScalarizeAccess(FixedVT, Idx))
+  if (!canScalarizeAccess(FixedVT, Idx, &I, AC))
     return false;
 
   InstructionCost OriginalCost = TTI.getMemoryOpCost(
       Instruction::Load, LI->getType(), Align(LI->getAlignment()),
       LI->getPointerAddressSpace());
   InstructionCost ScalarizedCost = 0;
 
   Instruction *LastCheckedInst = LI;

 INITIALIZE_PASS_END(VectorCombineLegacyPass, "vector-combine",
                     "Optimize scalar/vector ops", false, false)
 Pass *llvm::createVectorCombinePass() {
   return new VectorCombineLegacyPass();
 }
 
 PreservedAnalyses VectorCombinePass::run(Function &F,
                                          FunctionAnalysisManager &FAM) {
+  auto &AC = FAM.getResult<AssumptionAnalysis>(F);
   TargetTransformInfo &TTI = FAM.getResult<TargetIRAnalysis>(F);
   DominatorTree &DT = FAM.getResult<DominatorTreeAnalysis>(F);
   AAResults &AA = FAM.getResult<AAManager>(F);
-  VectorCombine Combiner(F, TTI, DT, AA);
+  VectorCombine Combiner(F, TTI, DT, AA, &AC);
   if (!Combiner.run())
     return PreservedAnalyses::all();
   PreservedAnalyses PA;
   PA.preserveSet<CFGAnalyses>();
   PA.preserve<GlobalsAA>();
   PA.preserve<AAManager>();
   PA.preserve<BasicAA>();
   return PA;
 }

llvm/test/Transforms/VectorCombine/AArch64/load-extractelement-scalarization.ll

entry:
   ret i32 %r
 }
 
 define i32 @load_extract_idx_var_i64_not_known_valid_by_assume_after_load(<4 x i32>* %x, i64 %idx) {
 ; CHECK-LABEL: @load_extract_idx_var_i64_not_known_valid_by_assume_after_load(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[IDX:%.*]], 4
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
-; CHECK-NEXT:    [[LV:%.*]] = load <4 x i32>, <4 x i32>* [[X:%.*]], align 16
-; CHECK-NEXT:    [[R:%.*]] = extractelement <4 x i32> [[LV]], i64 [[IDX]]
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds <4 x i32>, <4 x i32>* [[X:%.*]], i32 0, i64 [[IDX]]
+; CHECK-NEXT:    [[R:%.*]] = load i32, i32* [[TMP0]], align 1
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
 entry:
   %cmp = icmp ult i64 %idx, 4
   call void @llvm.assume(i1 %cmp)
   %lv = load <4 x i32>, <4 x i32>* %x
   %r = extractelement <4 x i32> %lv, i64 %idx
   ret i32 %r

llvm/test/Transforms/VectorCombine/load-insert-store.ll

entry:
   ret void
 }
 
 define void @insert_store_nonconst_index_known_valid_by_assume(<16 x i8>* %q, i8 zeroext %s, i32 %idx) {
 ; CHECK-LABEL: @insert_store_nonconst_index_known_valid_by_assume(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[IDX:%.*]], 4
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
-; CHECK-NEXT:    [[TMP0:%.*]] = load <16 x i8>, <16 x i8>* [[Q:%.*]], align 16
-; CHECK-NEXT:    [[VECINS:%.*]] = insertelement <16 x i8> [[TMP0]], i8 [[S:%.*]], i32 [[IDX]]
-; CHECK-NEXT:    store <16 x i8> [[VECINS]], <16 x i8>* [[Q]], align 16
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds <16 x i8>, <16 x i8>* [[Q:%.*]], i32 0, i32 [[IDX]]
+; CHECK-NEXT:    store i8 [[S:%.*]], i8* [[TMP0]], align 1
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %cmp = icmp ult i32 %idx, 4
   call void @llvm.assume(i1 %cmp)
   %0 = load <16 x i8>, <16 x i8>* %q
   %vecins = insertelement <16 x i8> %0, i8 %s, i32 %idx
   store <16 x i8> %vecins, <16 x i8>* %q
entry:
   %0 = load <16 x i8>, <16 x i8>* %q
   %vecins = insertelement <16 x i8> %0, i8 %s, i32 %idx
   store <16 x i8> %vecins, <16 x i8>* %q
   ret void
 }
 
 declare void @llvm.assume(i1)
 
 define void @insert_store_ptr_strip(<16 x i8>* %q, i8 zeroext %s) {
 ; CHECK-LABEL: @insert_store_ptr_strip(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[ADDR0:%.*]] = bitcast <16 x i8>* [[Q:%.*]] to <2 x i64>*
 ; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds <16 x i8>, <16 x i8>* [[Q]], i32 0, i32 3
 ; CHECK-NEXT:    store i8 [[S:%.*]], i8* [[TMP0]], align 1
 ; CHECK-NEXT:    ret void
 ;
 entry: