Add NVPTXLowerAlloca pass to convert alloca'ed memory to local address

[wengxt]

Summary

This is done by first adding two additional instructions to convert the
alloca returned address to local and convert it back to generic. Then
replace all uses of alloca instruction with the converted generic
address. Then we can rely NVPTXFavorNonGenericAddrSpace pass to combine
the generic addresscast and the corresponding Load, Store, Bitcast, GEP
Instruction together.

Patched by Xuetian Weng ([e-mail]).

[wengxt]

This will depends on D10482 to eliminate alloca across the addrspacecast boundary, otherwise there maybe performance regression because SROA will not work properly.

Though moving SROA before NVPTXLowerAlloca could also resolve the issue, but NVPTXLowerKernelArgs, NVPTXLowerAlloca, NVPTXFavorNonGenericAddrSpaces will merge someday, so it's better to fix this issue in SROA.

[jingyue]

lib/Target/NVPTX/NVPTXLowerAlloca.cpp

10

The comment here is imprecise. You're not adding "a cast to local"; you are adding a *pair* of addrspacecasts.

25

I'd mention the end result -- what the PTX looks like after this change, because that's what you are optimizing for.

65

Can this pass be a BasicBlockPass? Looks like all the changes are intra-BB.

66

Prefer

for (auto &B : F) {
  for (auto &I : B) {

unless you invalidate the iterators.

77

Similarly, I think there's a uses() that returns a range?

81

and BitCast. Can you comment in code why you only check these types? What would happen if you replace all uses?

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

52

Is this %SP + 0 removed due to your change? If so, add a CHECK-NOT in the test.

[jholewinski]

Other than the issues raised by Jingyue, this looks good to me! Thanks!

[jingyue]

feel free to throw in comments

[wengxt]

Try to address jingyue's comments.

For line 77 comment, the use iterator may be invalidated by changing
operands below.

for call-with-alloca-buffer.ll comment The %SP + 0 is not a critical
check, and it can appear in other position depending on some other
pass and instruction, so this check is removed and replaced by check
the dest register of cvta.to.local instruction.

The updated patch also move NVPTXLowerAlloca and NVPTXFavorNonGenericAddrSpaces
after SROA pass. So we don't need to depends on D10482 for now. This
can be changed back depending on how D10482 goes.

[jingyue]

LGTM with some minors

lib/Target/NVPTX/NVPTXLowerAlloca.cpp

15

; emits st.u32

just to show the difference from the optimized version

82

"Check Load, Store, GEP and BitCast Uses"

84

"For other types"

lib/Target/NVPTX/NVPTXTargetMachine.cpp

173

The first part of this comment still makes sense, right?

[wengxt]

Fix issues according to jingyue's comments

[jingyue]

lib/Target/NVPTX/NVPTXLowerAlloca.cpp

96

What are these "//continue"?

[wengxt]

sorry for mistake..

lib/Target/NVPTX/CMakeLists.txt

set(NVPTXCodeGen_sources
   NVPTXFrameLowering.cpp
   NVPTXGenericToNVVM.cpp
   NVPTXISelDAGToDAG.cpp
   NVPTXISelLowering.cpp
   NVPTXImageOptimizer.cpp
   NVPTXInstrInfo.cpp
   NVPTXLowerAggrCopies.cpp
   NVPTXLowerKernelArgs.cpp
+  NVPTXLowerAlloca.cpp
   NVPTXMCExpr.cpp
   NVPTXPrologEpilogPass.cpp
   NVPTXRegisterInfo.cpp
   NVPTXReplaceImageHandles.cpp
   NVPTXSubtarget.cpp
   NVPTXTargetMachine.cpp
   NVPTXTargetTransformInfo.cpp
   NVPTXUtilities.cpp
   NVVMReflect.cpp
   )
 
 add_llvm_target(NVPTXCodeGen ${NVPTXCodeGen_sources})
 
 add_subdirectory(TargetInfo)
 add_subdirectory(InstPrinter)
 add_subdirectory(MCTargetDesc)

lib/Target/NVPTX/NVPTX.h

 ModulePass *createGenericToNVVMPass();
 FunctionPass *createNVPTXFavorNonGenericAddrSpacesPass();
 ModulePass *createNVVMReflectPass();
 ModulePass *createNVVMReflectPass(const StringMap<int>& Mapping);
 MachineFunctionPass *createNVPTXPrologEpilogPass();
 MachineFunctionPass *createNVPTXReplaceImageHandlesPass();
 FunctionPass *createNVPTXImageOptimizerPass();
 FunctionPass *createNVPTXLowerKernelArgsPass(const NVPTXTargetMachine *TM);
+FunctionPass *createNVPTXLowerAllocaPass();
 
 bool isImageOrSamplerVal(const Value *, const Module *);
 
 extern Target TheNVPTXTarget32;
 extern Target TheNVPTXTarget64;
 
 namespace NVPTX {
 enum DrvInterface {

lib/Target/NVPTX/NVPTXLowerAlloca.cpp

+//===-- NVPTXLowerAlloca.cpp - Make alloca to use local memory =====--===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// For all alloca instructions, and add a cast to local address for
+// each of them. For example,
+//
+//   %A = alloca i32
+//   store i32 0, i32* %A
+//
+// will be transformed to
+//
+//   %A = alloca i32
+//   %Local = addrspacecast i32* %A to i32 addrspace(5)*
+//   %Generic = addrspacecast i32 addrspace(5)* %A to i32*
+//   store i32 0, i32 addrspace(5)* %Generic
+//
+// And we will rely on NVPTXFavorNonGenericAddrSpace to combine the last
+// two instructions.
+//
+//===----------------------------------------------------------------------===//
+
+#include "NVPTX.h"
+#include "NVPTXUtilities.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Type.h"
+#include "llvm/Pass.h"
+
+using namespace llvm;
+
+namespace llvm {
+void initializeNVPTXLowerAllocaPass(PassRegistry &);
+}
+
+namespace {
+class NVPTXLowerAlloca : public FunctionPass {
+  bool runOnFunction(Function &F) override;
+
+public:
+  static char ID; // Pass identification, replacement for typeid
+  NVPTXLowerAlloca() : FunctionPass(ID) {}
+  const char *getPassName() const override {
+    return "convert address space of alloca'ed memory to local";
+  }
+};
+} // namespace
+
+char NVPTXLowerAlloca::ID = 1;
+
+INITIALIZE_PASS(NVPTXLowerAlloca, "nvptx-lower-alloca",
+                "Lower Alloca", false, false)
+
+// =============================================================================
+// Main function for this pass.
+// =============================================================================
+bool NVPTXLowerAlloca::runOnFunction(Function &F) {
+  for (Function::iterator I = F.begin(), E = F.end(); I != E; ++I) {
+    for (BasicBlock::iterator BI = I->begin(), BE = I->end(); BI != BE;) {
+      if (auto allocaInst = dyn_cast<AllocaInst>(BI++)) {
+        auto PTy = dyn_cast<PointerType>(allocaInst->getType());
+        auto ETy = PTy->getElementType();
+        auto LocalAddrTy = PointerType::get(ETy, ADDRESS_SPACE_LOCAL);
+        auto NewASCToLocal = new AddrSpaceCastInst(allocaInst, LocalAddrTy, "");
+        auto GenericAddrTy = PointerType::get(ETy, ADDRESS_SPACE_GENERIC);
+        auto NewASCToGeneric = new AddrSpaceCastInst(NewASCToLocal,
+                                                     GenericAddrTy, "");
+        NewASCToLocal->insertAfter(allocaInst);
+        NewASCToGeneric->insertAfter(NewASCToLocal);
+        for (Value::use_iterator UI = allocaInst->use_begin(),
+                                 UE = allocaInst->use_end();
+             UI != UE; ) {
+          // Check Load, Store, GEP Uses on alloca and make them use the
+          // converted generic address.
+          const auto &AllocaUse = *UI++;
+          auto LI = dyn_cast<LoadInst>(AllocaUse.getUser());
+          if (LI && LI->getPointerOperand() == allocaInst) {
+            LI->setOperand(LI->getPointerOperandIndex(), NewASCToGeneric);
+            continue;
+          }
+          auto SI = dyn_cast<StoreInst>(AllocaUse.getUser());
+          if (SI && SI->getPointerOperand() == allocaInst) {
+            SI->setOperand(SI->getPointerOperandIndex(), NewASCToGeneric);
+          }
+          auto GI = dyn_cast<GetElementPtrInst>(AllocaUse.getUser());
+          if (GI && GI->getPointerOperand() == allocaInst) {
+            GI->setOperand(GI->getPointerOperandIndex(), NewASCToGeneric);
+          }
+          auto BI = dyn_cast<BitCastInst>(AllocaUse.getUser());
+          if (BI && BI->getOperand(0) == allocaInst) {
+            BI->setOperand(0, NewASCToGeneric);
+          }
+        }
+      }
+    }
+  }
+  return true;
+}
+
+FunctionPass *llvm::createNVPTXLowerAllocaPass() {
+  return new NVPTXLowerAlloca();
+}

lib/Target/NVPTX/NVPTXTargetMachine.cpp

 
 namespace llvm {
 void initializeNVVMReflectPass(PassRegistry&);
 void initializeGenericToNVVMPass(PassRegistry&);
 void initializeNVPTXAllocaHoistingPass(PassRegistry &);
 void initializeNVPTXAssignValidGlobalNamesPass(PassRegistry&);
 void initializeNVPTXFavorNonGenericAddrSpacesPass(PassRegistry &);
 void initializeNVPTXLowerKernelArgsPass(PassRegistry &);
+void initializeNVPTXLowerAllocaPass(PassRegistry &);
 }
 
 extern "C" void LLVMInitializeNVPTXTarget() {
   // Register the target.
   RegisterTargetMachine<NVPTXTargetMachine32> X(TheNVPTXTarget32);
   RegisterTargetMachine<NVPTXTargetMachine64> Y(TheNVPTXTarget64);
 
   // FIXME: This pass is really intended to be invoked during IR optimization,
   // but it's very NVPTX-specific.
   initializeNVVMReflectPass(*PassRegistry::getPassRegistry());
   initializeGenericToNVVMPass(*PassRegistry::getPassRegistry());
   initializeNVPTXAllocaHoistingPass(*PassRegistry::getPassRegistry());
   initializeNVPTXAssignValidGlobalNamesPass(*PassRegistry::getPassRegistry());
   initializeNVPTXFavorNonGenericAddrSpacesPass(
     *PassRegistry::getPassRegistry());
   initializeNVPTXLowerKernelArgsPass(*PassRegistry::getPassRegistry());
+  initializeNVPTXLowerAllocaPass(*PassRegistry::getPassRegistry());
 }
 
 static std::string computeDataLayout(bool is64Bit) {
   std::string Ret = "e";
 
   if (!is64Bit)
     Ret += "-p:32:32";
 
void NVPTXPassConfig::addIRPasses() {
   disablePass(&BranchFolderPassID);
   disablePass(&TailDuplicateID);
 
   addPass(createNVPTXImageOptimizerPass());
   TargetPassConfig::addIRPasses();
   addPass(createNVPTXAssignValidGlobalNamesPass());
   addPass(createGenericToNVVMPass());
   addPass(createNVPTXLowerKernelArgsPass(&getNVPTXTargetMachine()));
+  addPass(createNVPTXLowerAllocaPass());
   addPass(createNVPTXFavorNonGenericAddrSpacesPass());
   // NVPTXLowerKernelArgs emits alloca for byval parameters which can often
   // be eliminated by SROA. We do not run SROA right after NVPTXLowerKernelArgs
   // because we plan to merge NVPTXLowerKernelArgs and
   // NVPTXFavorNonGenericAddrSpaces into one pass.
   addPass(createSROAPass());
   // FavorNonGenericAddrSpaces shortcuts unnecessary addrspacecasts, and leave
   // them unused. We could remove dead code in an ad-hoc manner, but that
   // requires manual work and might be error-prone.
   addPass(createDeadCodeEliminationPass());

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

 ; CHECK: .local .align 4 .b8  __local_depot0[16]
 ; CHECK: mov.u64 %rd[[BUF_REG:[0-9]+]]
 ; CHECK: cvta.local.u64 %SP, %rd[[BUF_REG]]
 
 ; CHECK: ld.param.u64 %rd[[A_REG:[0-9]+]], [kernel_func_param_0]
 ; CHECK: cvta.to.global.u64 %rd[[A1_REG:[0-9]+]], %rd[[A_REG]]
 ; FIXME: casting A1_REG to A2_REG is unnecessary; A2_REG is essentially A_REG
 ; CHECK: cvta.global.u64 %rd[[A2_REG:[0-9]+]], %rd[[A1_REG]]
+; CHECK: cvta.local.u64 %rd[[SP_REG:[0-9]+]]
 ; CHECK: ld.global.f32 %f[[A0_REG:[0-9]+]], [%rd[[A1_REG]]]
-; CHECK: st.f32 [%SP+0], %f[[A0_REG]]
+; CHECK: st.local.f32 [{{%rd[0-9]+}}], %f[[A0_REG]]
 
   %0 = load float, float* %a, align 4
   %1 = bitcast [16 x i8]* %buf to float*
   store float %0, float* %1, align 4
   %arrayidx2 = getelementptr inbounds float, float* %a, i64 1
   %2 = load float, float* %arrayidx2, align 4
   %arrayidx3 = getelementptr inbounds [16 x i8], [16 x i8]* %buf, i64 0, i64 1
   %3 = bitcast i8* %arrayidx3 to float*
   store float %2, float* %3, align 4
   %arrayidx4 = getelementptr inbounds float, float* %a, i64 2
   %4 = load float, float* %arrayidx4, align 4
   %arrayidx5 = getelementptr inbounds [16 x i8], [16 x i8]* %buf, i64 0, i64 2
   %5 = bitcast i8* %arrayidx5 to float*
   store float %4, float* %5, align 4
   %arrayidx6 = getelementptr inbounds float, float* %a, i64 3
   %6 = load float, float* %arrayidx6, align 4
   %arrayidx7 = getelementptr inbounds [16 x i8], [16 x i8]* %buf, i64 0, i64 3
   %7 = bitcast i8* %arrayidx7 to float*
   store float %6, float* %7, align 4
 
-; CHECK: add.u64 %rd[[SP_REG:[0-9]+]], %SP, 0
 ; CHECK:        .param .b64 param0;
 ; CHECK-NEXT:   st.param.b64  [param0+0], %rd[[A2_REG]]
 ; CHECK-NEXT:   .param .b64 param1;
 ; CHECK-NEXT:   st.param.b64  [param1+0], %rd[[SP_REG]]
 ; CHECK-NEXT:   call.uni
 ; CHECK-NEXT:   callee,
 
   %arraydecay = getelementptr inbounds [16 x i8], [16 x i8]* %buf, i64 0, i64 0

test/CodeGen/NVPTX/lower-alloca.ll

+; RUN: opt < %s -S -nvptx-lower-alloca -nvptx-favor-non-generic -dce | FileCheck %s
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_35 | FileCheck %s --check-prefix PTX
+
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64"
+target triple = "nvptx64-unknown-unknown"
+
+define void @kernel() {
+; LABEL: @lower_alloca
+; PTX-LABEL: .visible .entry kernel(
+  %A = alloca i32
+; CHECK: addrspacecast i32* %A to i32 addrspace(5)*
+; CHECK: store i32 0, i32 addrspace(5)* {{%.+}}
+; PTX: st.local.u32 [{{%rd[0-9]+}}], {{%r[0-9]+}}
+  store i32 0, i32* %A
+  call void @callee(i32* %A)
+  ret void
+}
+
+declare void @callee(i32*)
+
+!nvvm.annotations = !{!0}
+!0 = !{void ()* @kernel, !"kernel", i32 1}