address comments

Author: AlexMaclean

llvm/lib/Target/NVPTX/NVPTXIncreaseAlignment.cpp

@@ -7,7 +7,7 @@
 //===----------------------------------------------------------------------===//
 //
 // A simple pass that looks at local memory arrays that are statically
-// sized and sets an appropriate alignment for them. This enables vectorization
+// sized and potentially increases their alignment. This enables vectorization
 // of loads/stores to these arrays if not explicitly specified by the client.
 //
 // TODO: Ideally we should do a bin-packing of local arrays to maximize
@@ -16,12 +16,15 @@
 //===----------------------------------------------------------------------===//
 
 #include "NVPTX.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/MathExtras.h"
+#include "llvm/Support/NVPTXAddrSpace.h"
 
 using namespace llvm;
 
@@ -30,16 +33,30 @@ static cl::opt<bool>
                            cl::init(false), cl::Hidden,
                            cl::desc("Use maximum alignment for local memory"));
 
-static constexpr Align MaxPTXArrayAlignment = Align::Constant<16>();
+namespace {
+struct NVPTXIncreaseLocalAlignment {
+  const Align MaxAlign;
+
+  NVPTXIncreaseLocalAlignment(const TargetTransformInfo &TTI)
+      : MaxAlign(TTI.getLoadStoreVecRegBitWidth(NVPTXAS::ADDRESS_SPACE_LOCAL)) {
+  }
+
+  bool run(Function &F);
+  bool updateAllocaAlignment(AllocaInst *Alloca, const DataLayout &DL);
+  Align getAggressiveArrayAlignment(unsigned ArraySize);
+  Align getConservativeArrayAlignment(unsigned ArraySize);
+};
+} // namespace
 
 /// Get the maximum useful alignment for an array. This is more likely to
 /// produce holes in the local memory.
 ///
 /// Choose an alignment large enough that the entire array could be loaded with
 /// a single vector load (if possible). Cap the alignment at
 /// MaxPTXArrayAlignment.
-static Align getAggressiveArrayAlignment(const unsigned ArraySize) {
-  return std::min(MaxPTXArrayAlignment, Align(PowerOf2Ceil(ArraySize)));
+Align NVPTXIncreaseLocalAlignment::getAggressiveArrayAlignment(
+    const unsigned ArraySize) {
+  return std::min(MaxAlign, Align(PowerOf2Ceil(ArraySize)));
 }
 
 /// Get the alignment of arrays that reduces the chances of leaving holes when
@@ -49,20 +66,18 @@ static Align getAggressiveArrayAlignment(const unsigned ArraySize) {
 /// Choose the largest alignment such that the array size is a multiple of the
 /// alignment. If all elements of the buffer are allocated in order of
 /// alignment (higher to lower) no holes will be left.
-static Align getConservativeArrayAlignment(const unsigned ArraySize) {
-  return commonAlignment(MaxPTXArrayAlignment, ArraySize);
+Align NVPTXIncreaseLocalAlignment::getConservativeArrayAlignment(
+    const unsigned ArraySize) {
+  return commonAlignment(MaxAlign, ArraySize);
 }
 
 /// Find a better alignment for local arrays
-static bool updateAllocaAlignment(const DataLayout &DL, AllocaInst *Alloca) {
+bool NVPTXIncreaseLocalAlignment::updateAllocaAlignment(AllocaInst *Alloca,
+                                                        const DataLayout &DL) {
   // Looking for statically sized local arrays
   if (!Alloca->isStaticAlloca())
     return false;
 
-  // For now, we only support array allocas
-  if (!(Alloca->isArrayAllocation() || Alloca->getAllocatedType()->isArrayTy()))
-    return false;
-
   const auto ArraySize = Alloca->getAllocationSize(DL);
   if (!(ArraySize && ArraySize->isFixed()))
     return false;
@@ -80,14 +95,14 @@ static bool updateAllocaAlignment(const DataLayout &DL, AllocaInst *Alloca) {
   return false;
 }
 
-static bool runSetLocalArrayAlignment(Function &F) {
+bool NVPTXIncreaseLocalAlignment::run(Function &F) {
   bool Changed = false;
-  const DataLayout &DL = F.getParent()->getDataLayout();
+  const auto &DL = F.getParent()->getDataLayout();
 
   BasicBlock &EntryBB = F.getEntryBlock();
   for (Instruction &I : EntryBB)
     if (AllocaInst *Alloca = dyn_cast<AllocaInst>(&I))
-      Changed |= updateAllocaAlignment(DL, Alloca);
+      Changed |= updateAllocaAlignment(Alloca, DL);
 
   return Changed;
 }
@@ -115,12 +130,15 @@ FunctionPass *llvm::createNVPTXIncreaseLocalAlignmentPass() {
 }
 
 bool NVPTXIncreaseLocalAlignmentLegacyPass::runOnFunction(Function &F) {
-  return runSetLocalArrayAlignment(F);
+  const auto &TTI = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
+  return NVPTXIncreaseLocalAlignment(TTI).run(F);
 }
 
 PreservedAnalyses
-NVPTXIncreaseLocalAlignmentPass::run(Function &F, FunctionAnalysisManager &AM) {
-  bool Changed = runSetLocalArrayAlignment(F);
+NVPTXIncreaseLocalAlignmentPass::run(Function &F,
+                                     FunctionAnalysisManager &FAM) {
+  const auto &TTI = FAM.getResult<TargetIRAnalysis>(F);
+  bool Changed = NVPTXIncreaseLocalAlignment(TTI).run(F);
 
   if (!Changed)
     return PreservedAnalyses::all();