LoadStoreVectorizer: Check TTI for vec reg bit width

[jlebar]

lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

471

Nit, since all we actually want is VecRegSize, consider making at least 'AS' a temporary:

Value *Ptr = LI->getPointerOperand();
unsigned VecRegSize = TTI.getLoadStoreVecRegBitWidth(Ptr, Ptr->getType()->getPointerAddressSpace());

This way readers have to carry one less local variable. Same elsewhere.

[arsenm]

r274322

include/llvm/LinkAllPasses.h

ForcePassLinking() {
       (void) llvm::createLintPass();
       (void) llvm::createSinkingPass();
       (void) llvm::createLowerAtomicPass();
       (void) llvm::createCorrelatedValuePropagationPass();
       (void) llvm::createMemDepPrinter();
       (void) llvm::createInstructionSimplifierPass();
       (void) llvm::createLoopVectorizePass();
       (void) llvm::createSLPVectorizerPass();
-      (void) llvm::createLoadStoreVectorizerPass(128);
+      (void) llvm::createLoadStoreVectorizerPass();
       (void) llvm::createBBVectorizePass();
       (void) llvm::createPartiallyInlineLibCallsPass();
       (void) llvm::createScalarizerPass();
       (void) llvm::createSeparateConstOffsetFromGEPPass();
       (void) llvm::createSpeculativeExecutionPass();
       (void) llvm::createSpeculativeExecutionIfHasBranchDivergencePass();
       (void) llvm::createRewriteSymbolsPass();
       (void) llvm::createStraightLineStrengthReducePass();

include/llvm/Transforms/Vectorize.h

bool vectorizeBasicBlock(Pass *P, BasicBlock &BB,
                          const VectorizeConfig &C = VectorizeConfig());
 
 
 //===----------------------------------------------------------------------===//
 //
 // LoadStoreVectorizer - Create vector loads and stores, but leave scalar
 // operations.
 //
-Pass *createLoadStoreVectorizerPass(unsigned VecRegSize);
+Pass *createLoadStoreVectorizerPass();
 
 } // End llvm namespace
 
 #endif

lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

class Vectorizer {
 
   typedef SmallVector<Value *, 8> ValueList;
   typedef MapVector<Value *, ValueList> ValueListMap;
 
   Function &F;
   AliasAnalysis &AA;
   DominatorTree &DT;
   ScalarEvolution &SE;
+  TargetTransformInfo &TTI;
   const DataLayout &DL;
   IRBuilder<> Builder;
   ValueListMap StoreRefs;
   ValueListMap LoadRefs;
-  unsigned VecRegSize;
 
 public:
   Vectorizer(Function &F, AliasAnalysis &AA, DominatorTree &DT,
-             ScalarEvolution &SE, unsigned VecRegSize)
-    : F(F), AA(AA), DT(DT), SE(SE), DL(F.getParent()->getDataLayout()),
-      Builder(SE.getContext()), VecRegSize(VecRegSize) {}
+             ScalarEvolution &SE, TargetTransformInfo &TTI)
+    : F(F), AA(AA), DT(DT), SE(SE), TTI(TTI),
+      DL(F.getParent()->getDataLayout()),
+      Builder(SE.getContext()) {}
 
   bool run();
 
 private:
   Value *getPointerOperand(Value *I);
 
   unsigned getAddressSpaceOperand(Value *I);
 
private:
 
   /// Vectorizes the store instructions in Chain.
   bool vectorizeStoreChain(ArrayRef<Value *> Chain);
 };
 
 class LoadStoreVectorizer : public FunctionPass {
 public:
   static char ID;
-  unsigned VecRegSize;
 
-  LoadStoreVectorizer(unsigned VecRegSize = 128) : FunctionPass(ID),
-                                                   VecRegSize(VecRegSize) {
+  LoadStoreVectorizer() : FunctionPass(ID) {
     initializeLoadStoreVectorizerPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnFunction(Function &F) override;
 
   const char *getPassName() const override {
     return "GPU Load and Store Vectorizer";
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AAResultsWrapperPass>();
     AU.addRequired<ScalarEvolutionWrapperPass>();
     AU.addRequired<DominatorTreeWrapperPass>();
+    AU.addRequired<TargetTransformInfoWrapperPass>();
     AU.setPreservesCFG();
   }
 };
 }
 
 INITIALIZE_PASS_BEGIN(LoadStoreVectorizer, DEBUG_TYPE,
                       "Vectorize load and Store instructions", false, false);
 INITIALIZE_PASS_DEPENDENCY(SCEVAAWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(GlobalsAAWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_END(LoadStoreVectorizer, DEBUG_TYPE,
                     "Vectorize load and store instructions", false, false);
 
 char LoadStoreVectorizer::ID = 0;
 
-Pass *llvm::createLoadStoreVectorizerPass(unsigned VecRegSize) {
-  return new LoadStoreVectorizer(VecRegSize);
+Pass *llvm::createLoadStoreVectorizerPass() {
+  return new LoadStoreVectorizer();
 }
 
 bool LoadStoreVectorizer::runOnFunction(Function &F) {
   AliasAnalysis &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
   DominatorTree &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
   ScalarEvolution &SE = getAnalysis<ScalarEvolutionWrapperPass>().getSE();
+  TargetTransformInfo &TTI
+    = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
 
   // Don't vectorize when the attribute NoImplicitFloat is used.
   if (F.hasFnAttribute(Attribute::NoImplicitFloat))
     return false;
 
-  Vectorizer V(F, AA, DT, SE, VecRegSize);
+  Vectorizer V(F, AA, DT, SE, TTI);
   return V.run();
 }
 
 // Vectorizer Implementation
 bool Vectorizer::run() {
   bool Changed = false;
 
   // Scan the blocks in the function in post order.
if (LoadInst *LI = dyn_cast<LoadInst>(it)) {
 
       if (!LI->isSimple())
         continue;
 
       Type *Ty = LI->getType();
       if (!VectorType::isValidElementType(Ty->getScalarType()))
         continue;
 
+      Value *Ptr = LI->getPointerOperand();
+      unsigned AS = Ptr->getType()->getPointerAddressSpace();
+      unsigned VecRegSize = TTI.getLoadStoreVecRegBitWidth(AS);
       // No point in looking at these if they're too big to vectorize.
       if (DL.getTypeSizeInBits(Ty) > VecRegSize / 2)
         continue;
 
       // Make sure all the users of a vector are constant-index extracts.
       if (isa<VectorType>(Ty) &&
           !std::all_of(LI->user_begin(), LI->user_end(), [LI] (const User *U) {
               const Instruction *UI = cast<Instruction>(U);
               return isa<ExtractElementInst>(UI) &&
                 isa<ConstantInt>(UI->getOperand(1));
             }))
         continue;
 
       // TODO: Target hook to filter types.
 
       // Save the load locations.
-      Value *Ptr = GetUnderlyingObject(LI->getPointerOperand(), DL);
-      LoadRefs[Ptr].push_back(LI);
+      Value *ObjPtr = GetUnderlyingObject(Ptr, DL);
+      LoadRefs[ObjPtr].push_back(LI);
 
     } else if (StoreInst *SI = dyn_cast<StoreInst>(it)) {
 
       if (!SI->isSimple())
         continue;
 
       Type *Ty = SI->getValueOperand()->getType();
       if (!VectorType::isValidElementType(Ty->getScalarType()))
         continue;
 
+      Value *Ptr = SI->getPointerOperand();
+      unsigned AS = Ptr->getType()->getPointerAddressSpace();
+      unsigned VecRegSize = TTI.getLoadStoreVecRegBitWidth(AS);
       if (DL.getTypeSizeInBits(Ty) > VecRegSize / 2)
         continue;
 
       if (isa<VectorType>(Ty) &&
           !std::all_of(SI->user_begin(), SI->user_end(), [SI] (const User *U) {
               const Instruction *UI = cast<Instruction>(U);
               return isa<ExtractElementInst>(UI) &&
                 isa<ConstantInt>(UI->getOperand(1));
             }))
         continue;
 
       // XXX: GPU-SPECIFIC CODE OMITTED. Here is where you want to omit
       // unsupported types or such.
 
       // Save store location.
-      Value *Ptr = GetUnderlyingObject(SI->getPointerOperand(), DL);
-      StoreRefs[Ptr].push_back(SI);
+      Value *ObjPtr = GetUnderlyingObject(Ptr, DL);
+      StoreRefs[ObjPtr].push_back(SI);
     }
   }
 }
 
 bool Vectorizer::vectorizeChains(ValueListMap &Map) {
   bool Changed = false;
 
   for (auto I = Map.begin(), E = Map.end(); I != E; ++I) {
bool Vectorizer::vectorizeInstructions(ArrayRef<Value *> Instrs) {
 
   return Changed;
 }
 
 bool Vectorizer::vectorizeStoreChain(ArrayRef<Value *> Chain) {
   StoreInst *S0 = cast<StoreInst>(Chain[0]);
   Type *StoreTy = S0->getValueOperand()->getType();
   unsigned Sz = DL.getTypeSizeInBits(StoreTy);
+  unsigned AS = S0->getPointerAddressSpace();
+  unsigned VecRegSize = TTI.getLoadStoreVecRegBitWidth(AS);
   unsigned VF = VecRegSize / Sz;
   unsigned ChainSize = Chain.size();
 
   if (!isPowerOf2_32(Sz) || VF < 2 || ChainSize < 2)
     return false;
 
   // Store size should be 1B, 2B or multiple of 4B.
   // TODO: Target hook for size constraint?
bool Vectorizer::vectorizeStoreChain(ArrayRef<Value *> Chain) {
   BasicBlock::iterator First = Boundaries.first;
   BasicBlock::iterator Last = Boundaries.second;
 
   if (!isVectorizable(Chain, First, Last))
     return false;
 
   // Set insert point.
   Builder.SetInsertPoint(&*Last);
-  unsigned AS = S0->getPointerAddressSpace();
 
   Value *Vec = UndefValue::get(VecTy);
 
   if (VecStoreTy) {
     unsigned VecWidth = VecStoreTy->getNumElements();
     for (unsigned i = 0; i < Chain.size(); ++i) {
       StoreInst *Store = cast<StoreInst>(Chain[i]);
       for (unsigned j = 0; j < VecStoreTy->getNumElements(); j++) {
bool Vectorizer::vectorizeStoreChain(ArrayRef<Value *> Chain) {
   ++NumVectorInstructions;
   return true;
 }
 
 bool Vectorizer::vectorizeLoadChain(ArrayRef<Value *> Chain) {
   LoadInst *L0 = cast<LoadInst>(Chain[0]);
   Type *LoadTy = L0->getType();
   unsigned Sz = DL.getTypeSizeInBits(LoadTy);
+  unsigned AS = L0->getPointerAddressSpace();
+  unsigned VecRegSize = TTI.getLoadStoreVecRegBitWidth(AS);
   unsigned VF = VecRegSize / Sz;
   unsigned ChainSize = Chain.size();
 
   if (!isPowerOf2_32(Sz) || VF < 2 || ChainSize < 2)
     return false;
 
   // Load size should be 1B, 2B or multiple of 4B.
   // TODO: Should size constraint be a target hook?
bool Vectorizer::vectorizeLoadChain(ArrayRef<Value *> Chain) {
   BasicBlock::iterator Last = Boundaries.second;
 
   if (!isVectorizable(Chain, First, Last))
     return false;
 
   // Set insert point.
   Builder.SetInsertPoint(&*Last);
 
-  unsigned AS = L0->getPointerAddressSpace();
   Value *Bitcast =
     Builder.CreateBitCast(L0->getPointerOperand(), VecTy->getPointerTo(AS));
 
   LoadInst *LI = cast<LoadInst>(Builder.CreateLoad(Bitcast));
   propagateMetadata(LI, Chain);
   LI->setAlignment(Alignment);
 
   SmallVector<Instruction *, 16> InstrsToReorder;

test/Transforms/LoadStoreVectorizer/AMDGPU/merge-stores-private.ll

+; RUN: opt -mtriple=amdgcn-amd-amdhsa -mattr=+max-private-element-size-4 -load-store-vectorizer -S -o - %s | FileCheck -check-prefix=ELT4 -check-prefix=ALL %s
+; RUN: opt -mtriple=amdgcn-amd-amdhsa -mattr=+max-private-element-size-8 -load-store-vectorizer -S -o - %s | FileCheck -check-prefix=ELT8 -check-prefix=ALL %s
+; RUN: opt -mtriple=amdgcn-amd-amdhsa -mattr=+max-private-element-size-16 -load-store-vectorizer -S -o - %s | FileCheck -check-prefix=ELT16 -check-prefix=ALL %s
+
+; ALL-LABEL: @merge_private_store_4_vector_elts_loads_v4i32
+; ELT4: store i32
+; ELT4: store i32
+; ELT4: store i32
+; ELT4: store i32
+
+; ELT8: store <2 x i32>
+; ELT8: store <2 x i32>
+
+; ELT16: store <4 x i32>
+define void @merge_private_store_4_vector_elts_loads_v4i32(i32* %out) #0 {
+  %out.gep.1 = getelementptr i32, i32* %out, i32 1
+  %out.gep.2 = getelementptr i32, i32* %out, i32 2
+  %out.gep.3 = getelementptr i32, i32* %out, i32 3
+
+  store i32 9, i32* %out
+  store i32 1, i32* %out.gep.1
+  store i32 23, i32* %out.gep.2
+  store i32 19, i32* %out.gep.3
+  ret void
+}
+
+; ALL-LABEL: @merge_private_store_4_vector_elts_loads_v4i8(
+; ALL: store <4 x i8>
+define void @merge_private_store_4_vector_elts_loads_v4i8(i8* %out) #0 {
+  %out.gep.1 = getelementptr i8, i8* %out, i32 1
+  %out.gep.2 = getelementptr i8, i8* %out, i32 2
+  %out.gep.3 = getelementptr i8, i8* %out, i32 3
+
+  store i8 9, i8* %out
+  store i8 1, i8* %out.gep.1
+  store i8 23, i8* %out.gep.2
+  store i8 19, i8* %out.gep.3
+  ret void
+}
+
+; ALL-LABEL: @merge_private_store_4_vector_elts_loads_v2i16(
+; ALL: store <2 x i16>
+define void @merge_private_store_4_vector_elts_loads_v2i16(i16* %out) #0 {
+  %out.gep.1 = getelementptr i16, i16* %out, i32 1
+
+  store i16 9, i16* %out
+  store i16 12, i16* %out.gep.1
+  ret void
+}
+
+attributes #0 = { nounwind }

test/Transforms/LoadStoreVectorizer/AMDGPU/merge-stores.ll

define void @merge_local_store_2_constants_i32_align_2(i32 addrspace(3)* %out) #0 {
   %out.gep.1 = getelementptr i32, i32 addrspace(3)* %out, i32 1
 
   store i32 123, i32 addrspace(3)* %out.gep.1, align 2
   store i32 456, i32 addrspace(3)* %out, align 2
   ret void
 }
 
 ; CHECK-LABEL: @merge_local_store_4_constants_i32
-; CHECK: store <4 x i32> <i32 1234, i32 123, i32 456, i32 333>, <4 x i32> addrspace(3)*
+; CHECK: store <2 x i32> <i32 456, i32 333>, <2 x i32> addrspace(3)*
+; CHECK: store <2 x i32> <i32 1234, i32 123>, <2 x i32> addrspace(3)*
 define void @merge_local_store_4_constants_i32(i32 addrspace(3)* %out) #0 {
   %out.gep.1 = getelementptr i32, i32 addrspace(3)* %out, i32 1
   %out.gep.2 = getelementptr i32, i32 addrspace(3)* %out, i32 2
   %out.gep.3 = getelementptr i32, i32 addrspace(3)* %out, i32 3
 
   store i32 123, i32 addrspace(3)* %out.gep.1
   store i32 456, i32 addrspace(3)* %out.gep.2
   store i32 333, i32 addrspace(3)* %out.gep.3