[NVPTX] Add NVPTXIncreaseAligmentPass to improve vectorization

llvm/lib/Target/NVPTX/CMakeLists.txt

@@ -26,6 +26,7 @@ set(NVPTXCodeGen_sources
   NVPTXISelLowering.cpp
   NVPTXLowerAggrCopies.cpp
   NVPTXLowerAlloca.cpp
+  NVPTXIncreaseAlignment.cpp
   NVPTXLowerArgs.cpp
   NVPTXLowerUnreachable.cpp
   NVPTXMCExpr.cpp

llvm/lib/Target/NVPTX/NVPTX.h

@@ -55,6 +55,7 @@ FunctionPass *createNVPTXTagInvariantLoadsPass();
 MachineFunctionPass *createNVPTXPeephole();
 MachineFunctionPass *createNVPTXProxyRegErasurePass();
 MachineFunctionPass *createNVPTXForwardParamsPass();
+FunctionPass *createNVPTXIncreaseLocalAlignmentPass();
 
 void initializeNVVMReflectLegacyPassPass(PassRegistry &);
 void initializeGenericToNVVMLegacyPassPass(PassRegistry &);
@@ -77,6 +78,7 @@ void initializeNVPTXExternalAAWrapperPass(PassRegistry &);
 void initializeNVPTXPeepholePass(PassRegistry &);
 void initializeNVPTXTagInvariantLoadLegacyPassPass(PassRegistry &);
 void initializeNVPTXPrologEpilogPassPass(PassRegistry &);
+void initializeNVPTXIncreaseLocalAlignmentLegacyPassPass(PassRegistry &);
 
 struct NVVMIntrRangePass : PassInfoMixin<NVVMIntrRangePass> {
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
@@ -112,6 +114,11 @@ struct NVPTXTagInvariantLoadsPass : PassInfoMixin<NVPTXTagInvariantLoadsPass> {
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 
+struct NVPTXIncreaseLocalAlignmentPass
+    : PassInfoMixin<NVPTXIncreaseLocalAlignmentPass> {
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+
 namespace NVPTX {
 enum DrvInterface {
   NVCL,

llvm/lib/Target/NVPTX/NVPTXIncreaseAlignment.cpp

@@ -0,0 +1,160 @@
+//===-- NVPTXIncreaseAlignment.cpp - Increase alignment for local arrays --===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// A simple pass that looks at local memory allocas that are statically
+// sized and potentially increases their alignment. This enables vectorization
+// of loads/stores to these allocas if not explicitly specified by the client.
+//
+// TODO: Ideally we should do a bin-packing of local allocas to maximize
+// alignments while minimizing holes.
+//
+//===----------------------------------------------------------------------===//
+
+#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;
+
+static cl::opt<unsigned> MinLocalArrayAlignment(
+    "nvptx-ensure-minimum-local-alignment", cl::init(16), cl::Hidden,
+    cl::desc(
+        "Ensure local memory objects are at least this aligned (default 16)"));
+
+static Align getMaxLocalArrayAlignment(const TargetTransformInfo &TTI) {
+  const unsigned MaxBitWidth =
+      TTI.getLoadStoreVecRegBitWidth(NVPTXAS::ADDRESS_SPACE_LOCAL);
+  return Align(MaxBitWidth / 8);
+}
+
+namespace {
+struct NVPTXIncreaseLocalAlignment {
+  const Align MaxUsableAlign;
+
+  NVPTXIncreaseLocalAlignment(const TargetTransformInfo &TTI)
+      : MaxUsableAlign(getMaxLocalArrayAlignment(TTI)) {}
+
+  bool run(Function &F);
+  bool updateAllocaAlignment(AllocaInst *Alloca, const DataLayout &DL);
+  Align getMaxUsefulArrayAlignment(unsigned ArraySize);
+  Align getMaxSafeLocalAlignment(unsigned ArraySize);
+};
+} // namespace
+
+/// Get the maximum useful alignment for an allocation. This is more likely to
+/// produce holes in the local memory.
+///
+/// Choose an alignment large enough that the entire alloca could be loaded
+/// with a single vector load (if possible). Cap the alignment at
+/// MinLocalArrayAlignment and MaxUsableAlign.
+Align NVPTXIncreaseLocalAlignment::getMaxUsefulArrayAlignment(
+    const unsigned ArraySize) {
+  const Align UpperLimit =
+      std::min(MaxUsableAlign, Align(MinLocalArrayAlignment));
+  return std::min(UpperLimit, Align(PowerOf2Ceil(ArraySize)));
+}
+
+/// Get the alignment of allocas that reduces the chances of leaving holes when
+/// they are allocated within a contiguous memory buffer (like the stack).
+/// Holes are still possible before and after the allocation.
+///
+/// Choose the largest alignment such that the allocation 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.
+Align NVPTXIncreaseLocalAlignment::getMaxSafeLocalAlignment(
+    const unsigned ArraySize) {
+  return commonAlignment(MaxUsableAlign, ArraySize);
+}
+
+/// Find a better alignment for local allocas.
+bool NVPTXIncreaseLocalAlignment::updateAllocaAlignment(AllocaInst *Alloca,
+                                                        const DataLayout &DL) {
+  if (!Alloca->isStaticAlloca())
+    return false;
+
+  const auto ArraySize = Alloca->getAllocationSize(DL);
+  if (!(ArraySize && ArraySize->isFixed()))
+    return false;
+
+  const auto ArraySizeValue = ArraySize->getFixedValue();
+  if (ArraySizeValue == 0)
+    return false;
+
+  const Align NewAlignment =
+      std::max(getMaxSafeLocalAlignment(ArraySizeValue),
+               getMaxUsefulArrayAlignment(ArraySizeValue));
+
+  if (NewAlignment > Alloca->getAlign()) {
+    Alloca->setAlignment(NewAlignment);
+    return true;
+  }
+
+  return false;
+}
+
+bool NVPTXIncreaseLocalAlignment::run(Function &F) {
+  bool Changed = false;
+  const auto &DL = F.getParent()->getDataLayout();
+
+  BasicBlock &EntryBB = F.getEntryBlock();
+  for (Instruction &I : EntryBB)
+    if (AllocaInst *Alloca = dyn_cast<AllocaInst>(&I))
+      Changed |= updateAllocaAlignment(Alloca, DL);
+
+  return Changed;
+}
+
+namespace {
+struct NVPTXIncreaseLocalAlignmentLegacyPass : public FunctionPass {
+  static char ID;
+  NVPTXIncreaseLocalAlignmentLegacyPass() : FunctionPass(ID) {}
+
+  bool runOnFunction(Function &F) override;
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<TargetTransformInfoWrapperPass>();
+  }
+  StringRef getPassName() const override {
+    return "NVPTX Increase Local Alignment";
+  }
+};
+} // namespace
+
+char NVPTXIncreaseLocalAlignmentLegacyPass::ID = 0;
+INITIALIZE_PASS(NVPTXIncreaseLocalAlignmentLegacyPass,
+                "nvptx-increase-local-alignment",
+                "Increase alignment for statically sized allocas", false, false)
+
+FunctionPass *llvm::createNVPTXIncreaseLocalAlignmentPass() {
+  return new NVPTXIncreaseLocalAlignmentLegacyPass();
+}
+
+bool NVPTXIncreaseLocalAlignmentLegacyPass::runOnFunction(Function &F) {
+  const auto &TTI = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
+  return NVPTXIncreaseLocalAlignment(TTI).run(F);
+}
+
+PreservedAnalyses
+NVPTXIncreaseLocalAlignmentPass::run(Function &F,
+                                     FunctionAnalysisManager &FAM) {
+  const auto &TTI = FAM.getResult<TargetIRAnalysis>(F);
+  bool Changed = NVPTXIncreaseLocalAlignment(TTI).run(F);
+
+  if (!Changed)
+    return PreservedAnalyses::all();
+
+  PreservedAnalyses PA;
+  PA.preserveSet<CFGAnalyses>();
+  return PA;
+}

llvm/lib/Target/NVPTX/NVPTXPassRegistry.def

@@ -40,4 +40,5 @@ FUNCTION_PASS("nvvm-intr-range", NVVMIntrRangePass())
 FUNCTION_PASS("nvptx-copy-byval-args", NVPTXCopyByValArgsPass())
 FUNCTION_PASS("nvptx-lower-args", NVPTXLowerArgsPass(*this))
 FUNCTION_PASS("nvptx-tag-invariant-loads", NVPTXTagInvariantLoadsPass())
+FUNCTION_PASS("nvptx-increase-local-alignment", NVPTXIncreaseLocalAlignmentPass())
 #undef FUNCTION_PASS

llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp

@@ -393,6 +393,8 @@ void NVPTXPassConfig::addIRPasses() {
   // but EarlyCSE can do neither of them.
   if (getOptLevel() != CodeGenOptLevel::None) {
     addEarlyCSEOrGVNPass();
+    // Increase alignment for local arrays to improve vectorization.
+    addPass(createNVPTXIncreaseLocalAlignmentPass());
     if (!DisableLoadStoreVectorizer)
       addPass(createLoadStoreVectorizerPass());
     addPass(createSROAPass());

llvm/test/CodeGen/NVPTX/call-with-alloca-buffer.ll

@@ -20,7 +20,7 @@ define ptx_kernel void @kernel_func(ptr %a) {
 entry:
   %buf = alloca [16 x i8], align 4
 
-; CHECK: .local .align 4 .b8 	__local_depot0[16]
+; CHECK: .local .align 16 .b8 	__local_depot0[16]
 ; CHECK: mov.b64 %SPL
 
 ; CHECK: ld.param.b64 %rd[[A_REG:[0-9]+]], [kernel_func_param_0]

llvm/test/CodeGen/NVPTX/increase-local-align.ll

@@ -0,0 +1,102 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=nvptx-increase-local-alignment -nvptx-ensure-minimum-local-alignment=1 < %s | FileCheck %s --check-prefixes=COMMON,MIN-1
+; RUN: opt -S -passes=nvptx-increase-local-alignment -nvptx-ensure-minimum-local-alignment=8 < %s | FileCheck %s --check-prefixes=COMMON,MIN-8
+; RUN: opt -S -passes=nvptx-increase-local-alignment -nvptx-ensure-minimum-local-alignment=16 < %s | FileCheck %s --check-prefixes=COMMON,MIN-16
+target triple = "nvptx64-nvidia-cuda"
+
+define void @test1() {
+; COMMON-LABEL: define void @test1() {
+; COMMON-NEXT:    [[A:%.*]] = alloca i8, align 1
+; COMMON-NEXT:    ret void
+;
+  %a = alloca i8, align 1
+  ret void
+}
+
+define void @test2() {
+; MIN-1-LABEL: define void @test2() {
+; MIN-1-NEXT:    [[A:%.*]] = alloca [63 x i8], align 1
+; MIN-1-NEXT:    ret void
+;
+; MIN-8-LABEL: define void @test2() {
+; MIN-8-NEXT:    [[A:%.*]] = alloca [63 x i8], align 8
+; MIN-8-NEXT:    ret void
+;
+; MIN-16-LABEL: define void @test2() {
+; MIN-16-NEXT:    [[A:%.*]] = alloca [63 x i8], align 16
+; MIN-16-NEXT:    ret void
+;
+  %a = alloca [63 x i8], align 1
+  ret void
+}
+
+define void @test3() {
+; COMMON-LABEL: define void @test3() {
+; COMMON-NEXT:    [[A:%.*]] = alloca [64 x i8], align 16
+; COMMON-NEXT:    ret void
+;
+  %a = alloca [64 x i8], align 1
+  ret void
+}
+
+define void @test4() {
+; MIN-1-LABEL: define void @test4() {
+; MIN-1-NEXT:    [[A:%.*]] = alloca i8, i32 63, align 1
+; MIN-1-NEXT:    ret void
+;
+; MIN-8-LABEL: define void @test4() {
+; MIN-8-NEXT:    [[A:%.*]] = alloca i8, i32 63, align 8
+; MIN-8-NEXT:    ret void
+;
+; MIN-16-LABEL: define void @test4() {
+; MIN-16-NEXT:    [[A:%.*]] = alloca i8, i32 63, align 16
+; MIN-16-NEXT:    ret void
+;
+  %a = alloca i8, i32 63, align 1
+  ret void
+}
+
+define void @test5() {
+; COMMON-LABEL: define void @test5() {
+; COMMON-NEXT:    [[A:%.*]] = alloca i8, i32 64, align 16
+; COMMON-NEXT:    ret void
+;
+  %a = alloca i8, i32 64, align 1
+  ret void
+}
+
+define void @test6() {
+; COMMON-LABEL: define void @test6() {
+; COMMON-NEXT:    [[A:%.*]] = alloca i8, align 32
+; COMMON-NEXT:    ret void
+;
+  %a = alloca i8, align 32
+  ret void
+}
+
+define void @test7() {
+; COMMON-LABEL: define void @test7() {
+; COMMON-NEXT:    [[A:%.*]] = alloca i32, align 4
+; COMMON-NEXT:    ret void
+;
+  %a = alloca i32, align 2
+  ret void
+}
+
+define void @test8() {
+; COMMON-LABEL: define void @test8() {
+; COMMON-NEXT:    [[A:%.*]] = alloca [2 x i32], align 8
+; COMMON-NEXT:    ret void
+;
+  %a = alloca [2 x i32], align 2
+  ret void
+}
+
+define void @test9() {
+; COMMON-LABEL: define void @test9() {
+; COMMON-NEXT:    [[A:%.*]] = alloca [0 x i32], align 1
+; COMMON-NEXT:    ret void
+;
+  %a = alloca [0 x i32], align 1
+  ret void
+}

llvm/test/CodeGen/NVPTX/local-stack-frame.ll

@@ -94,7 +94,7 @@ declare void @bar(ptr %a)
 define void @foo3(i32 %a) {
 ; PTX32-LABEL: foo3(
 ; PTX32:       {
-; PTX32-NEXT:    .local .align 4 .b8 __local_depot2[12];
+; PTX32-NEXT:    .local .align 16 .b8 __local_depot2[16];
 ; PTX32-NEXT:    .reg .b32 %SP;
 ; PTX32-NEXT:    .reg .b32 %SPL;
 ; PTX32-NEXT:    .reg .b32 %r<6>;
@@ -110,7 +110,7 @@ define void @foo3(i32 %a) {
 ;
 ; PTX64-LABEL: foo3(
 ; PTX64:       {
-; PTX64-NEXT:    .local .align 4 .b8 __local_depot2[12];
+; PTX64-NEXT:    .local .align 16 .b8 __local_depot2[16];
 ; PTX64-NEXT:    .reg .b64 %SP;
 ; PTX64-NEXT:    .reg .b64 %SPL;
 ; PTX64-NEXT:    .reg .b32 %r<2>;