[SYCL] Local Accessor Support for CUDA

Provides the LocalAccessorToSharedMemory compiler pass required
for supporting SYCL local accessors in CUDA.

Contributors
Alexander Johnston <[email protected]>
David Wood <[email protected]>

Signed-off-by: Alexander Johnston <[email protected]>

Author: Alexander Johnston

clang/lib/CodeGen/BackendUtil.cpp

@@ -842,9 +842,6 @@ void EmitAssemblyHelper::EmitAssembly(BackendAction Action,
   PerFunctionPasses.add(
       createTargetTransformInfoWrapperPass(getTargetIRAnalysis()));
 
-  if (LangOpts.SYCLIsDevice)
-    PerFunctionPasses.add(createSYCLLowerWGScopePass());
-
   CreatePasses(PerModulePasses, PerFunctionPasses);
 
   legacy::PassManager CodeGenPasses;

clang/lib/CodeGen/CodeGenAction.cpp

@@ -10,6 +10,7 @@
 #include "CodeGenModule.h"
 #include "CoverageMappingGen.h"
 #include "MacroPPCallbacks.h"
+#include "SYCLLowerIR/LowerWGScope.h"
 #include "clang/AST/ASTConsumer.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/DeclCXX.h"
@@ -33,6 +34,7 @@
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/LLVMRemarkStreamer.h"
+#include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/Linker/Linker.h"
@@ -326,6 +328,17 @@ namespace clang {
           CodeGenOpts.getProfileUse() != CodeGenOptions::ProfileNone)
         Ctx.setDiagnosticsHotnessRequested(true);
 
+      // The parallel_for_work_group legalization pass can emit calls to
+      // builtins function. Definitions of those builtins can be provided in
+      // LinkModule. We force the pass to legalize the code before the link
+      // happens.
+      if (LangOpts.SYCLIsDevice) {
+        PrettyStackTraceString CrashInfo("Pre-linking SYCL passes");
+        legacy::PassManager PreLinkingSyclPasses;
+        PreLinkingSyclPasses.add(createSYCLLowerWGScopePass());
+        PreLinkingSyclPasses.run(*getModule());
+      }
+
       // Link each LinkModule into our module.
       if (LinkInModules())
         return;

clang/lib/CodeGen/SYCLLowerIR/LowerWGScope.cpp

@@ -124,7 +124,8 @@ class SYCLLowerWGScopeLegacyPass : public FunctionPass {
       return false;
 
     FunctionAnalysisManager FAM;
-    auto PA = Impl.run(F, FAM);
+    auto TT = llvm::Triple(F.getParent()->getTargetTriple());
+    auto PA = Impl.run(F, TT, FAM);
     return !PA.areAllPreserved();
   }
 
@@ -188,8 +189,8 @@ enum class MemorySemantics : unsigned {
   ImageMemory = 0x800,
 };
 
-Instruction *genWGBarrier(Instruction &Before);
-Value *genLinearLocalID(Instruction &Before);
+Instruction *genWGBarrier(Instruction &Before, const Triple &TT);
+Value *genLinearLocalID(Instruction &Before, const Triple &TT);
 GlobalVariable *createWGLocalVariable(Module &M, Type *T, const Twine &Name);
 } // namespace spirv
 
@@ -263,8 +264,9 @@ static bool mayHaveSideEffects(const Instruction *I) {
 //
 static void guardBlockWithIsLeaderCheck(BasicBlock *IfBB, BasicBlock *TrueBB,
                                         BasicBlock *MergeBB,
-                                        const DebugLoc &DbgLoc) {
-  Value *LinearLocalID = spirv::genLinearLocalID(*IfBB->getTerminator());
+                                        const DebugLoc &DbgLoc,
+                                        const Triple &TT) {
+  Value *LinearLocalID = spirv::genLinearLocalID(*IfBB->getTerminator(), TT);
   auto *Ty = LinearLocalID->getType();
   Value *Zero = Constant::getNullValue(Ty);
   IRBuilder<> Builder(IfBB->getContext());
@@ -341,7 +343,7 @@ using InstrRange = std::pair<Instruction *, Instruction *>;
 //   ...
 //   B
 //   ... USE2(%I1_new) ...
-static void tformRange(const InstrRange &R) {
+static void tformRange(const InstrRange &R, const Triple &TT) {
   // Instructions seen between the first and the last
   SmallPtrSet<Instruction *, 16> Seen;
   Instruction *FirstSE = R.first;
@@ -360,15 +362,15 @@ static void tformRange(const InstrRange &R) {
 
   // 1) insert the first "is work group leader" test (at the first split) for
   //     the worker WIs to detour the side effects instructions
-  guardBlockWithIsLeaderCheck(BBa, LeaderBB, BBb, FirstSE->getDebugLoc());
+  guardBlockWithIsLeaderCheck(BBa, LeaderBB, BBb, FirstSE->getDebugLoc(), TT);
 
   // 2) "Share" the output values of the instructions in the range
   for (auto *I : Seen)
     shareOutputViaLocalMem(*I, *BBa, *BBb, Seen);
 
   // 3) Insert work group barrier so that workers further read valid data
   //    (before the materialization reads inserted at step 2)
-  spirv::genWGBarrier(BBb->front());
+  spirv::genWGBarrier(BBb->front(), TT);
 }
 
 namespace {
@@ -443,13 +445,13 @@ static void copyBetweenPrivateAndShadow(Value *L, GlobalVariable *Shadow,
 //
 static void materializeLocalsInWIScopeBlocksImpl(
     const DenseMap<BasicBlock *, std::unique_ptr<LocalsSet>> &BB2MatLocals,
-    const DenseMap<AllocaInst *, GlobalVariable *> &Local2Shadow) {
+    const DenseMap<AllocaInst *, GlobalVariable *> &Local2Shadow, const Triple &TT) {
   for (auto &P : BB2MatLocals) {
     // generate LeaderBB and private<->shadow copies in proper BBs
     BasicBlock *LeaderBB = P.first;
     BasicBlock *BB = LeaderBB->splitBasicBlock(&LeaderBB->front(), "LeaderMat");
     // Add a barrier to the original block:
-    Instruction *At = spirv::genWGBarrier(*BB->getFirstNonPHI())->getNextNode();
+    Instruction *At = spirv::genWGBarrier(*BB->getFirstNonPHI(), TT)->getNextNode();
 
     for (AllocaInst *L : *P.second.get()) {
       auto MapEntry = Local2Shadow.find(L);
@@ -472,7 +474,7 @@ static void materializeLocalsInWIScopeBlocksImpl(
     BasicBlock *TestBB =
         LeaderBB->splitBasicBlock(&LeaderBB->front(), "TestMat");
     std::swap(TestBB, LeaderBB);
-    guardBlockWithIsLeaderCheck(TestBB, LeaderBB, BB, At->getDebugLoc());
+    guardBlockWithIsLeaderCheck(TestBB, LeaderBB, BB, At->getDebugLoc(), TT);
   }
 }
 
@@ -536,7 +538,8 @@ static bool localMustBeMaterialized(const AllocaInst *L, const BasicBlock &BB) {
 //
 void materializeLocalsInWIScopeBlocks(
     SmallPtrSetImpl<AllocaInst *> &Locals,
-    SmallPtrSetImpl<BasicBlock *> &WIScopeBBs) {
+    SmallPtrSetImpl<BasicBlock *> &WIScopeBBs,
+    const Triple &TT) {
   // maps local variable to its "shadow" workgroup-shared global:
   DenseMap<AllocaInst *, GlobalVariable *> Local2Shadow;
   // records which locals must be materialized at the beginning of a block:
@@ -567,7 +570,7 @@ void materializeLocalsInWIScopeBlocks(
     }
   }
   // perform the materialization
-  materializeLocalsInWIScopeBlocksImpl(BB2MatLocals, Local2Shadow);
+  materializeLocalsInWIScopeBlocksImpl(BB2MatLocals, Local2Shadow, TT);
 }
 
 #ifndef NDEBUG
@@ -680,7 +683,7 @@ static void fixupPrivateMemoryPFWILambdaCaptures(CallInst *PFWICall) {
 // Go through "byval" parameters which are passed as AS(0) pointers
 // and: (1) create local shadows for them (2) and initialize them from the
 // leader's copy and (3) replace usages with pointer to the shadow
-static void shareByValParams(Function &F) {
+static void shareByValParams(Function &F, const Triple &TT) {
   // split
   BasicBlock *EntryBB = &F.getEntryBlock();
   BasicBlock *LeaderBB = EntryBB->splitBasicBlock(&EntryBB->front(), "leader");
@@ -689,7 +692,7 @@ static void shareByValParams(Function &F) {
   // 1) rewire the above basic blocks so that LeaderBB is executed only for the
   // leader workitem
   guardBlockWithIsLeaderCheck(EntryBB, LeaderBB, MergeBB,
-                              EntryBB->back().getDebugLoc());
+                              EntryBB->back().getDebugLoc(), TT);
   Instruction &At = LeaderBB->back();
 
   for (auto &Arg : F.args()) {
@@ -715,10 +718,11 @@ static void shareByValParams(Function &F) {
                                 true /*private->shadow*/);
   }
   // 5) make sure workers use up-to-date shared values written by the leader
-  spirv::genWGBarrier(MergeBB->front());
+  spirv::genWGBarrier(MergeBB->front(), TT);
 }
 
 PreservedAnalyses SYCLLowerWGScopePass::run(Function &F,
+                                            const llvm::Triple &TT,
                                             FunctionAnalysisManager &FAM) {
   if (!F.getMetadata(WG_SCOPE_MD))
     return PreservedAnalyses::all();
@@ -796,7 +800,7 @@ PreservedAnalyses SYCLLowerWGScopePass::run(Function &F,
 
   // Perform the transformation
   for (auto &R : Ranges) {
-    tformRange(R);
+    tformRange(R, TT);
     Changed = true;
   }
   // There can be allocas not corresponding to any variable declared in user
@@ -813,14 +817,14 @@ PreservedAnalyses SYCLLowerWGScopePass::run(Function &F,
     WIScopeBBs.insert(I->getParent());
 
   // Now materialize the locals:
-  materializeLocalsInWIScopeBlocks(Allocas, WIScopeBBs);
+  materializeLocalsInWIScopeBlocks(Allocas, WIScopeBBs, TT);
 
   // Fixup captured addresses of private_memory isntances in current WI
   for (auto *PFWICall : PFWICalls)
     fixupPrivateMemoryPFWILambdaCaptures(PFWICall);
 
   // Finally, create shadows for and replace usages of byval pointer params
-  shareByValParams(F);
+  shareByValParams(F, TT);
 
 #ifndef NDEBUG
   if (HaveChanges && Debug > 0)
@@ -866,37 +870,74 @@ GlobalVariable *spirv::createWGLocalVariable(Module &M, Type *T,
 // Must correspond to the code in
 // llvm-spirv/lib/SPIRV/OCL20ToSPIRV.cpp
 // OCL20ToSPIRV::transWorkItemBuiltinsToVariables()
-Value *spirv::genLinearLocalID(Instruction &Before) {
+Value *spirv::genLinearLocalID(Instruction &Before, const Triple &TT) {
   Module &M = *Before.getModule();
-  StringRef Name = "__spirv_BuiltInLocalInvocationIndex";
-  GlobalVariable *G = M.getGlobalVariable(Name);
-
-  if (!G) {
-    Type *T = getSizeTTy(M);
-    G = new GlobalVariable(M,                              // module
-                           T,                              // type
-                           true,                           // isConstant
-                           GlobalValue::ExternalLinkage,   // Linkage
-                           nullptr,                        // Initializer
-                           Name,                           // Name
-                           nullptr,                        // InsertBefore
-                           GlobalVariable::NotThreadLocal, // ThreadLocalMode
-                           // TODO 'Input' crashes CPU Back-End
-                           // asUInt(spirv::AddrSpace::Input) // AddressSpace
-                           asUInt(spirv::AddrSpace::Global) // AddressSpace
-    );
-    unsigned Align = M.getDataLayout().getPreferredAlignment(G);
-    G->setAlignment(MaybeAlign(Align));
+  if (TT.isNVPTX()) {
+    LLVMContext &Ctx = Before.getContext();
+    Type *RetTy = getSizeTTy(M);
+
+    IRBuilder<> Bld(Ctx);
+    Bld.SetInsertPoint(&Before);
+
+#define CREATE_CALLEE(NAME, FN_NAME) \
+  FunctionCallee FnCallee##NAME = M.getOrInsertFunction(FN_NAME, RetTy); \
+  assert(FnCallee##NAME && "spirv intrinsic creation failed"); \
+  auto NAME = Bld.CreateCall(FnCallee##NAME, {});
+
+      CREATE_CALLEE(LocalInvocationId_X, "_Z27__spirv_LocalInvocationId_xv");
+      CREATE_CALLEE(LocalInvocationId_Y, "_Z27__spirv_LocalInvocationId_yv");
+      CREATE_CALLEE(LocalInvocationId_Z, "_Z27__spirv_LocalInvocationId_zv");
+      CREATE_CALLEE(WorkgroupSize_Y, "_Z23__spirv_WorkgroupSize_yv");
+      CREATE_CALLEE(WorkgroupSize_Z, "_Z23__spirv_WorkgroupSize_zv");
+
+#undef CREATE_CALLEE
+
+    // 1:   ((__spirv_WorkgroupSize_y() * __spirv_WorkgroupSize_z())
+    // 2:    * __spirv_LocalInvocationId_x())
+    // 3: + (__spirv_WorkgroupSize_z() * __spirv_LocalInvocationId_y())
+    // 4: + (__spirv_LocalInvocationId_z())
+    return Bld.CreateAdd(
+      Bld.CreateAdd(
+        Bld.CreateMul(
+          Bld.CreateMul(WorkgroupSize_Y, WorkgroupSize_Z), // 1
+          LocalInvocationId_X), // 2
+        Bld.CreateMul(WorkgroupSize_Z, LocalInvocationId_Y)), // 3
+      LocalInvocationId_Z); // 4
+  } else {
+    StringRef Name = "__spirv_BuiltInLocalInvocationIndex";
+    GlobalVariable *G = M.getGlobalVariable(Name);
+
+    if (!G) {
+      Type *T = getSizeTTy(M);
+      G = new GlobalVariable(M,                              // module
+                             T,                              // type
+                             true,                           // isConstant
+                             GlobalValue::ExternalLinkage,   // Linkage
+                             nullptr,                        // Initializer
+                             Name,                           // Name
+                             nullptr,                        // InsertBefore
+                             GlobalVariable::NotThreadLocal, // ThreadLocalMode
+                             // TODO 'Input' crashes CPU Back-End
+                             // asUInt(spirv::AddrSpace::Input) // AddressSpace
+                             asUInt(spirv::AddrSpace::Global) // AddressSpace
+      );
+      unsigned Align = M.getDataLayout().getPreferredAlignment(G);
+      G->setAlignment(Align);
+    }
+    Value *Res = new LoadInst(G, "", &Before);
+    return Res;
   }
-  Value *Res = new LoadInst(G, "", &Before);
-  return Res;
 }
 
 // extern void __spirv_ControlBarrier(Scope Execution, Scope Memory,
 //  uint32_t Semantics) noexcept;
-Instruction *spirv::genWGBarrier(Instruction &Before) {
+Instruction *spirv::genWGBarrier(Instruction &Before, const Triple &TT) {
   Module &M = *Before.getModule();
-  StringRef Name = "__spirv_ControlBarrier";
+  StringRef Name;
+  if (TT.isNVPTX())
+    Name = "_Z22__spirv_ControlBarrierN5__spv5ScopeES0_j";
+  else
+    Name = "__spirv_ControlBarrier";
   LLVMContext &Ctx = Before.getContext();
   Type *ScopeTy = Type::getInt32Ty(Ctx);
   Type *SemanticsTy = Type::getInt32Ty(Ctx);

clang/lib/CodeGen/SYCLLowerIR/LowerWGScope.h

@@ -21,7 +21,7 @@ namespace llvm {
 /// execution model semantics - this code must be executed once per work group.
 class SYCLLowerWGScopePass : public PassInfoMixin<SYCLLowerWGScopePass> {
 public:
-  PreservedAnalyses run(Function &F, FunctionAnalysisManager &);
+  PreservedAnalyses run(Function &F, const Triple &TT, FunctionAnalysisManager &);
 };
 
 FunctionPass *createSYCLLowerWGScopePass();

llvm/lib/Target/NVPTX/CMakeLists.txt

@@ -33,6 +33,7 @@ set(NVPTXCodeGen_sources
   NVVMIntrRange.cpp
   NVVMReflect.cpp
   NVPTXProxyRegErasure.cpp
+  SYCL/LocalAccessorToSharedMemory.cpp
   )
 
 add_llvm_target(NVPTXCodeGen ${NVPTXCodeGen_sources})

llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp

@@ -17,6 +17,7 @@
 #include "NVPTXTargetObjectFile.h"
 #include "NVPTXTargetTransformInfo.h"
 #include "TargetInfo/NVPTXTargetInfo.h"
+#include "SYCL/LocalAccessorToSharedMemory.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
@@ -70,6 +71,8 @@ void initializeNVPTXLowerArgsPass(PassRegistry &);
 void initializeNVPTXLowerAllocaPass(PassRegistry &);
 void initializeNVPTXProxyRegErasurePass(PassRegistry &);
 
+void initializeLocalAccessorToSharedMemoryPass(PassRegistry &);
+
 } // end namespace llvm
 
 extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeNVPTXTarget() {
@@ -89,6 +92,9 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeNVPTXTarget() {
   initializeNVPTXLowerAllocaPass(PR);
   initializeNVPTXLowerAggrCopiesPass(PR);
   initializeNVPTXProxyRegErasurePass(PR);
+
+  // SYCL-specific passes, needed here to be available to `opt`.
+  initializeLocalAccessorToSharedMemoryPass(PR);
 }
 
 static std::string computeDataLayout(bool is64Bit, bool UseShortPointers) {
@@ -266,6 +272,11 @@ void NVPTXPassConfig::addIRPasses() {
   const NVPTXSubtarget &ST = *getTM<NVPTXTargetMachine>().getSubtargetImpl();
   addPass(createNVVMReflectPass(ST.getSmVersion()));
 
+  if (getTM<NVPTXTargetMachine>().getTargetTriple().getOS() == Triple::CUDA &&
+      getTM<NVPTXTargetMachine>().getTargetTriple().getEnvironment() == Triple::SYCLDevice) {
+    addPass(createLocalAccessorToSharedMemoryPass());
+  }
+
   if (getOptLevel() != CodeGenOpt::None)
     addPass(createNVPTXImageOptimizerPass());
   addPass(createNVPTXAssignValidGlobalNamesPass());