[OpenMP] Enable the lowering of implicitly shared variables in OpenMP GPU-offloaded target regions to the GPU shared memory

[gtbercea]

Summary

This patch is part of the development effort to add support in the current OpenMP GPU offloading implementation for implicitly sharing variables between a target region executed by the team master thread and the worker threads within that team.

This patch is the second of three required for successfully performing the implicit sharing of master thread variables with the worker threads within a team:
-Patch D38976 extends the CLANG code generation with code that handles shared variables.
-Patch (coming soon) extends the functionality of libomptarget to maintain a list of references to shared variables.

This patch adds a shared memory stack to the prolog of the kernel function representing the device offloaded OpenMP target region. The new passes along with the changes to existing ones, ensure that any OpenMP variable which needs to be shared across several threads will be allocated in this new stack, in the shared memory of the device. This patch covers the case of sharing variables from the master thread to the worker threads:

#pragma omp target
{
   // master thread only
   int v;
   #pragma omp parallel
   {
      // worker threads
      // use v
   }
}

[tra]

Please add tests for the cases where such local->shaed conversion should and should not happen.
I would appreciate if you could add details on what exactly your passes are supposed to move to shared memory.

Considering that device-side code tends to be heavily inlined, it may be prudent to add an option to control the total size of shared memory we allow to be used for this purpose.

In case your passes are not executed (or didn't move anything to shared memory), is there any impact on the generated PTX. I.e. can ptxas successfully optimize unused shared memory away?

If the code intentionally wants to allocate something in local memory, would the allocation ever be moved to shared memory by your pass? If so, how would I prevent that?

lib/Target/NVPTX/NVPTXAsmPrinter.cpp

1753

Nit: the name should end with S as the L in SPL was for 'local' address space. which then gets converted to generic AS. In your case it will be in shared space, hence S would be more appropriate.

lib/Target/NVPTX/NVPTXAssignValidGlobalNames.cpp

68 ↗

(On Diff #119210)

The name cleanup changes in this file should probably be committed by themselves as they have nothing to do with the rest of the patch.

lib/Target/NVPTX/NVPTXFunctionDataSharing.cpp

10

Please add details about what the pass is supposed to do.

[gtbercea]

Eliminate variable and function name clean-up. That has been moved into a separate patch: D39005

[gtbercea]

Add regression tests and allow for shared memory lowering to be disabled at function level.

[gtbercea]

ping

[gtbercea]

@tra @hfinkel

[gtbercea]

ping

[guansong]

Here is a question, do we require that the alloca size to be compile time constant?

[hfinkel]

lib/Target/NVPTX/NVPTXAsmPrinter.cpp

1737

Line too long.

lib/Target/NVPTX/NVPTXFrameLowering.cpp

71

In other places in this patch you refer explicitly to OpenMP, so it probably makes sense to say "the OpenMP runtime" here as well (but just saying "the runtime" seems potentially confusing).

85

Line too long.

lib/Target/NVPTX/NVPTXLowerSharedFrameIndicesPass.cpp

13

Can you be more specific? I believe that we fixed PEI to handle virtual registers, so if that's the only motivation, can we use the regular PEI now?

lib/Target/NVPTX/NVPTXRegisterInfo.cpp

134

Line too long.

lib/Target/NVPTX/NVPTXUtilities.cpp

322

Can't you use PointerMayBeCaptured (include/llvm/Analysis/CaptureTracking.h) instead of this function? If so, please do.

[gtbercea]

Use LLVM function for checking if pointer is stored.

[gtbercea]

ping

[tra]

Dotting the 'i's on the questions that were not replied to directly.

In D38978#899205, @tra wrote:

Considering that device-side code tends to be heavily inlined, it may be prudent to add an option to control the total size of shared memory we allow to be used for this purpose.

I'm still curious to hear what do you plan to do when your depot use grows beyond certain limit. At the very least there's the physical limit on shared memory size. Shared memory use also affects how many threads can be launched which has large impact on performance. IMO having some sort of user-controllable threshold would be very desirable.

In case your passes are not executed (or didn't move anything to shared memory), is there any impact on the generated PTX. I.e. can ptxas successfully optimize unused shared memory away?

This may have been addressed by the no-shared-depot.ll test. It would be nice to add few comments in the tests explaining what they do.

If the code intentionally wants to allocate something in local memory, would the allocation ever be moved to shared memory by your pass? If so, how would I prevent that?

AFAICT this functionality only applies to functions with has-nvptx-shared-depot attribute. Works for me.

lib/Target/NVPTX/NVPTXFunctionDataSharing.cpp

99

Nit: return false would match the intent better.

lib/Target/NVPTX/NVPTXRegisterInfo.td

75

Line too long.

test/CodeGen/NVPTX/insert-shared-depot.ll

5–6

You could put common checks under the same label (e.g. CHECK) and run tests with -check-prefixes=PTX32,CHECK

30

'LABEL' is not a check-prefix and @linsert_shared_depot is not this function's name, so I'm puzzled what this line is supposed to do. Did you intend <prefix>-LABEL: @kernel ?

This appears in all the test cases in the patch.

[gtbercea]

Address comments.

[gtbercea]

In D38978#967485, @tra wrote:

Dotting the 'i's on the questions that were not replied to directly.

In D38978#899205, @tra wrote:

Considering that device-side code tends to be heavily inlined, it may be prudent to add an option to control the total size of shared memory we allow to be used for this purpose.

I'm still curious to hear what do you plan to do when your depot use grows beyond certain limit. At the very least there's the physical limit on shared memory size. Shared memory use also affects how many threads can be launched which has large impact on performance. IMO having some sort of user-controllable threshold would be very desirable.

When shared memory isn't enough to hold the shared depot, global memory will be used instead. That is a scheme which will be covered by a future patch.

In case your passes are not executed (or didn't move anything to shared memory), is there any impact on the generated PTX. I.e. can ptxas successfully optimize unused shared memory away?

This may have been addressed by the no-shared-depot.ll test. It would be nice to add few comments in the tests explaining what they do.

Done.

If the code intentionally wants to allocate something in local memory, would the allocation ever be moved to shared memory by your pass? If so, how would I prevent that?

AFAICT this functionality only applies to functions with has-nvptx-shared-depot attribute. Works for me.

That's right.

test/CodeGen/NVPTX/insert-shared-depot.ll

30

This is modeled after the lower-alloca.ll test which has a similar label. The label is always equal to the name of the test file. In this particular case there is a typo, it should be "insert_shared_depot" not "linsert_shared_depot"

[tra]

In D38978#968222, @gtbercea wrote:

I'm still curious to hear what do you plan to do when your depot use grows beyond certain limit. At the very least there's the physical limit on shared memory size. Shared memory use also affects how many threads can be launched which has large impact on performance. IMO having some sort of user-controllable threshold would be very desirable.

When shared memory isn't enough to hold the shared depot, global memory will be used instead. That is a scheme which will be covered by a future patch.

Good luck with that. IMO if your kernel requires all shared memory available per multiprocessor, you are almost guaranteed suboptimal performance because you will not have enough threads running -- neither for peak compute, nor to hide global memory access latency. My bet that you will eventually end up limiting shared memory use to a fairly small fraction of it.

Given that impact is limited to explicitly annotated functions only, this lack of tune-ability is OK with me for now. I'd add a TODO item somewhere to describe that tuning specific limits is WIP.

test/CodeGen/NVPTX/insert-shared-depot.ll

30

This is modeled after the lower-alloca.ll test which has a similar label.

lower-alloca.ll indeed has the same problem.

The label is always equal to the name of the test file.

I don't think FileCheck has such a feature. Nor do I see anything matching this description in the FileCheck documentation. Nor does it work. See below.

In this particular case there is a typo, it should be "insert_shared_depot" not "linsert_shared_depot"

The line does not check *anything* right now. In this test FileCheck only pays attention to lines that have CHECK or PTX64/PTX32. This line contains neither and is ignored. You can do an experiment -- replace the line with ; LABEL: this should never match and run the test.

I've tried that on lower-alloca.ll and the test, as expected, passes regardless of the nonsense I put after the LABEL:.

[gtbercea]

In D38978#968565, @tra wrote:

In D38978#968222, @gtbercea wrote:

I'm still curious to hear what do you plan to do when your depot use grows beyond certain limit. At the very least there's the physical limit on shared memory size. Shared memory use also affects how many threads can be launched which has large impact on performance. IMO having some sort of user-controllable threshold would be very desirable.

When shared memory isn't enough to hold the shared depot, global memory will be used instead. That is a scheme which will be covered by a future patch.

Good luck with that. IMO if your kernel requires all shared memory available per multiprocessor, you are almost guaranteed suboptimal performance because you will not have enough threads running -- neither for peak compute, nor to hide global memory access latency. My bet that you will eventually end up limiting shared memory use to a fairly small fraction of it.

I completely agree, this scheme will be efficient only when modest amounts of shared memory are required, for larger memory footprints, a global memory scheme will be used instead.

Given that impact is limited to explicitly annotated functions only, this lack of tune-ability is OK with me for now. I'd add a TODO item somewhere to describe that tuning specific limits is WIP.

I'll choose a sensible default for the cut-off point/condition and make it tune-able by the user once we have the global memory scheme in place.

[gtbercea]

Remove LABEL from tests and add TODO comment for shared memory limit.

[Hahnfeld]

Not my area of expertise

[gtbercea]

ping

[gtbercea]

Alternative solution was implemented.

include/llvm/CodeGen/TargetPassConfig.h

protected:
   virtual bool addPreISel() {
     return true;
   }
 
   /// addMachineSSAOptimization - Add standard passes that optimize machine
   /// instructions in SSA form.
   virtual void addMachineSSAOptimization();
 
+  /// Add passes that lower variables to a
+  /// particular memory type.
+  virtual void addMachineSSALowering() {}
+
   /// Add passes that optimize instruction level parallelism for out-of-order
   /// targets. These passes are run while the machine code is still in SSA
   /// form, so they can use MachineTraceMetrics to control their heuristics.
   ///
   /// All passes added here should preserve the MachineDominatorTree,
   /// MachineLoopInfo, and MachineTraceMetrics analyses.
   virtual bool addILPOpts() {
     return false;

lib/CodeGen/TargetPassConfig.cpp

void TargetPassConfig::addMachinePasses() {
 
   // Expand pseudo-instructions emitted by ISel.
   addPass(&ExpandISelPseudosID);
 
   // Add passes that optimize machine instructions in SSA form.
   if (getOptLevel() != CodeGenOpt::None) {
     addMachineSSAOptimization();
   } else {
+    // Ensure lowering to the appropriate memroy type occurs even when no
+    // optimizations are enabled. This type of lowering is required for
+    // correctness by the NVPTX backend.
+    addMachineSSALowering();
+
     // If the target requests it, assign local variables to stack slots relative
     // to one another and simplify frame index references where possible.
     addPass(&LocalStackSlotAllocationID, false);
   }
 
   if (TM->Options.EnableIPRA)
     addPass(createRegUsageInfoPropPass());
 

lib/Target/NVPTX/CMakeLists.txt

set(NVPTXCodeGen_sources
   NVPTXImageOptimizer.cpp
   NVPTXInstrInfo.cpp
   NVPTXLowerAggrCopies.cpp
   NVPTXLowerArgs.cpp
   NVPTXLowerAlloca.cpp
   NVPTXPeephole.cpp
   NVPTXMCExpr.cpp
   NVPTXPrologEpilogPass.cpp
+  NVPTXLowerSharedFrameIndicesPass.cpp
   NVPTXRegisterInfo.cpp
   NVPTXReplaceImageHandles.cpp
   NVPTXSubtarget.cpp
   NVPTXTargetMachine.cpp
   NVPTXTargetTransformInfo.cpp
+  NVPTXFunctionDataSharing.cpp
   NVPTXUtilities.cpp
   NVVMIntrRange.cpp
   NVVMReflect.cpp
   )
 
 add_llvm_target(NVPTXCodeGen ${NVPTXCodeGen_sources})
 
 add_subdirectory(TargetInfo)
 add_subdirectory(InstPrinter)
 add_subdirectory(MCTargetDesc)

lib/Target/NVPTX/NVPTX.h

 
 FunctionPass *createNVPTXISelDag(NVPTXTargetMachine &TM,
                                  llvm::CodeGenOpt::Level OptLevel);
 ModulePass *createNVPTXAssignValidGlobalNamesPass();
 ModulePass *createGenericToNVVMPass();
 FunctionPass *createNVVMIntrRangePass(unsigned int SmVersion);
 FunctionPass *createNVVMReflectPass();
 MachineFunctionPass *createNVPTXPrologEpilogPass();
+MachineFunctionPass *createNVPTXLowerSharedFrameIndicesPass();
 MachineFunctionPass *createNVPTXReplaceImageHandlesPass();
 FunctionPass *createNVPTXImageOptimizerPass();
 FunctionPass *createNVPTXLowerArgsPass(const NVPTXTargetMachine *TM);
 BasicBlockPass *createNVPTXLowerAllocaPass();
+FunctionPass *createNVPTXFunctionDataSharingPass(const NVPTXTargetMachine *TM);
 MachineFunctionPass *createNVPTXPeephole();
 
 Target &getTheNVPTXTarget32();
 Target &getTheNVPTXTarget64();
 
 namespace NVPTX {
 enum DrvInterface {
   NVCL,

lib/Target/NVPTX/NVPTXAsmPrinter.cpp

 #include <sstream>
 #include <string>
 #include <utility>
 #include <vector>
 
 using namespace llvm;
 
 #define DEPOTNAME "__local_depot"
+#define SHARED_DEPOTNAME "__shared_depot"
 
 static cl::opt<bool>
 EmitLineNumbers("nvptx-emit-line-numbers", cl::Hidden,
                 cl::desc("NVPTX Specific: Emit Line numbers even without -G"),
                 cl::init(true));
 
 static cl::opt<bool>
 InterleaveSrc("nvptx-emit-src", cl::ZeroOrMore, cl::Hidden,
void NVPTXAsmPrinter::setAndEmitFunctionVirtualRegisters(
     const MachineFunction &MF) {
   SmallString<128> Str;
   raw_svector_ostream O(Str);
 
   // Map the global virtual register number to a register class specific
   // virtual register number starting from 1 with that class.
   const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
   //unsigned numRegClasses = TRI->getNumRegClasses();
+  bool IsKernelFunction = isKernelFunction(MF.getFunction());
+
+  bool GenerateSharedDepot =
+      MF.getFunction().hasFnAttribute("has-nvptx-shared-depot");
 
   // Emit the Fake Stack Object
   const MachineFrameInfo &MFI = MF.getFrameInfo();
   int NumBytes = (int) MFI.getStackSize();
   if (NumBytes) {
     O << "\t.local .align " << MFI.getMaxAlignment() << " .b8 \t" << DEPOTNAME
       << getFunctionNumber() << "[" << NumBytes << "];\n";
+    if (IsKernelFunction && GenerateSharedDepot) {
+      O << "\t.shared .align " << MFI.getMaxAlignment()
+        << " .b8 \t" << SHARED_DEPOTNAME << getFunctionNumber()
+        << "[" << NumBytes << "];\n";
+    }
     if (static_cast<const NVPTXTargetMachine &>(MF.getTarget()).is64Bit()) {
       O << "\t.reg .b64 \t%SP;\n";
       O << "\t.reg .b64 \t%SPL;\n";
+      if (IsKernelFunction && GenerateSharedDepot) {
+        O << "\t.reg .b64 \t%SPS;\n";
+        O << "\t.reg .b64 \t%SPSH;\n";
+      }
     } else {
       O << "\t.reg .b32 \t%SP;\n";
       O << "\t.reg .b32 \t%SPL;\n";
+      if (IsKernelFunction && GenerateSharedDepot) {
+        O << "\t.reg .b32 \t%SPS;\n";
+        O << "\t.reg .b32 \t%SPSH;\n";
+      }
     }
   }
 
   // Go through all virtual registers to establish the mapping between the
   // global virtual
   // register number and the per class virtual register number.
   // We use the per class virtual register number in the ptx output.
   unsigned int numVRs = MRI->getNumVirtRegs();

 void NVPTXAsmPrinter::printOperand(const MachineInstr *MI, int opNum,
                                    raw_ostream &O, const char *Modifier) {
   const MachineOperand &MO = MI->getOperand(opNum);
   switch (MO.getType()) {
   case MachineOperand::MO_Register:
     if (TargetRegisterInfo::isPhysicalRegister(MO.getReg())) {
       if (MO.getReg() == NVPTX::VRDepot)
         O << DEPOTNAME << getFunctionNumber();
+      else if (MO.getReg() == NVPTX::VRSharedDepot)
+        O << SHARED_DEPOTNAME << getFunctionNumber();
       else
         O << NVPTXInstPrinter::getRegisterName(MO.getReg());
     } else {
       emitVirtualRegister(MO.getReg(), O);
     }
     return;
 
   case MachineOperand::MO_Immediate:

lib/Target/NVPTX/NVPTXFrameLowering.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "NVPTXFrameLowering.h"
 #include "NVPTX.h"
 #include "NVPTXRegisterInfo.h"
 #include "NVPTXSubtarget.h"
 #include "NVPTXTargetMachine.h"
+#include "NVPTXUtilities.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/MC/MachineLocation.h"
 
 using namespace llvm;
if (!MR.use_empty(NVPTX::VRFrame)) {
       MI = BuildMI(MBB, MI, dl,
                    MF.getSubtarget().getInstrInfo()->get(CvtaLocalOpcode),
                    NVPTX::VRFrame)
                .addReg(NVPTX::VRFrameLocal);
     }
     BuildMI(MBB, MI, dl, MF.getSubtarget().getInstrInfo()->get(MovDepotOpcode),
             NVPTX::VRFrameLocal)
         .addImm(MF.getFunctionNumber());
+
+    bool SharedStackPointerInit =
+        MF.getFunction().hasFnAttribute("has-nvptx-shared-depot");
+
+    // Only emit a shared depot for the main kernel function.
+    // The other device functions need to get a handle on this shared depot
+    // by interacting with a runtime library.
+    //
+    // Clang can trigger this at any point via the has-nvptx-shared-depot function
+    // attribute.
+    //
+    // Currently this situation arises as a consequence of OpenMP semantics.
+    // The interaction in that case invovles the OpenMP runtime.
+    if (isKernelFunction(MF.getFunction()) && SharedStackPointerInit) {
+      // Emits
+      //   mov %SHSPL, %shared_depot;
+      //   cvta.shared %SHSP, %SHSPL;
+      // For the time being just emit it even if it's not used.
+      unsigned CvtaSharedOpcode =
+          Is64Bit ? NVPTX::cvta_shared_yes_64 : NVPTX::cvta_shared_yes;
+      unsigned MovSharedDepotOpcode =
+          Is64Bit ? NVPTX::MOV_SHARED_DEPOT_ADDR_64 : NVPTX::MOV_SHARED_DEPOT_ADDR;
+      MI = BuildMI(MBB, MI, dl,
+                   MF.getSubtarget().getInstrInfo()->get(CvtaSharedOpcode),
+                   NVPTX::VRShared)
+               .addReg(NVPTX::VRFrameShared);
+      BuildMI(MBB, MI, dl,
+              MF.getSubtarget().getInstrInfo()->get(MovSharedDepotOpcode),
+              NVPTX::VRFrameShared)
+          .addImm(MF.getFunctionNumber());
+    }
   }
 }
 
 void NVPTXFrameLowering::emitEpilogue(MachineFunction &MF,
                                       MachineBasicBlock &MBB) const {}
 
 // This function eliminates ADJCALLSTACKDOWN,
 // ADJCALLSTACKUP pseudo instructions
 MachineBasicBlock::iterator NVPTXFrameLowering::eliminateCallFramePseudoInstr(
     MachineFunction &MF, MachineBasicBlock &MBB,
     MachineBasicBlock::iterator I) const {
   // Simply discard ADJCALLSTACKDOWN,
   // ADJCALLSTACKUP instructions.
   return MBB.erase(I);
 }

lib/Target/NVPTX/NVPTXFunctionDataSharing.h

+//===--- NVPTXFrameLowering.h - Define frame lowering for NVPTX -*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+//
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_TARGET_NVPTX_NVPTXFUNCTIONDATASHARING_H
+#define LLVM_LIB_TARGET_NVPTX_NVPTXFUNCTIONDATASHARING_H
+
+namespace llvm {
+
+class NVPTXFunctionDataSharing : public FunctionPass {
+  bool runOnFunction(Function &F) override;
+  bool runOnKernelFunction(Function &F);
+  bool runOnDeviceFunction(Function &F);
+
+public:
+  static char ID; // Pass identification, replacement for typeid
+  NVPTXFunctionDataSharing(const NVPTXTargetMachine *TM = nullptr)
+      : FunctionPass(ID), TM(TM) {}
+  StringRef getPassName() const override {
+    return "Function level data sharing pass.";
+  }
+
+private:
+  const NVPTXTargetMachine *TM;
+};
+} // End llvm namespace
+
+#endif
 No newline at end of file

lib/Target/NVPTX/NVPTXFunctionDataSharing.cpp

+//===-- FunctionDataSharing.cpp - Mark pointers as shared -----------------===//
+//
+//                     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, add a pair of cast to shared address for
+// each of them. For example,
+//
+//   %A = alloca i32
+//   store i32 0, i32* %A ; emits st.u32
+//
+// will be transformed to
+//
+//   %A = alloca i32
+//   %Local = addrspacecast i32* %A to i32 addrspace(3)*
+//   %Shared = addrspacecast i32 addrspace(3)* %A to i32*
+//   store i32 0, i32 addrspace(5)* %Shared ; emits st.shared.u32
+//
+// And we will rely on NVPTXInferAddressSpaces to combine the last two
+// instructions.
+//
+// This pass is invoked for -O0 only.
+//
+//===----------------------------------------------------------------------===//
+
+#include "NVPTX.h"
+#include "NVPTXUtilities.h"
+#include "NVPTXTargetMachine.h"
+#include "llvm/Analysis/ValueTracking.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Type.h"
+#include "llvm/Pass.h"
+
+using namespace llvm;
+
+namespace llvm {
+void initializeNVPTXFunctionDataSharingPass(PassRegistry &);
+}
+
+namespace {
+class NVPTXFunctionDataSharing : public FunctionPass {
+  bool runOnFunction(Function &F) override;
+  bool runOnKernelFunction(Function &F);
+  bool runOnDeviceFunction(Function &F);
+
+public:
+  static char ID; // Pass identification, replacement for typeid
+  NVPTXFunctionDataSharing(const NVPTXTargetMachine *TM = nullptr)
+      : FunctionPass(ID) {}
+  StringRef getPassName() const override {
+    return "Function level data sharing pass.";
+  }
+};
+} // namespace
+
+char NVPTXFunctionDataSharing::ID = 1;
+
+INITIALIZE_PASS(NVPTXFunctionDataSharing, "nvptx-function-data-sharing",
+                "Function Data Sharing (NVPTX)", false, false)
+
+static void markPointerAsShared(Value *Ptr) {
+  if (Ptr->getType()->getPointerAddressSpace() == ADDRESS_SPACE_SHARED)
+    return;
+
+  // Deciding where to emit the addrspacecast pair.
+  // Insert right after Ptr if Ptr is an instruction.
+  BasicBlock::iterator InsertPt =
+      std::next(cast<Instruction>(Ptr)->getIterator());
+  assert(InsertPt != InsertPt->getParent()->end() &&
+         "We don't call this function with Ptr being a terminator.");
+
+  auto *PtrInShared = new AddrSpaceCastInst(
+      Ptr, PointerType::get(Ptr->getType()->getPointerElementType(),
+                            ADDRESS_SPACE_SHARED),
+      Ptr->getName(), &*InsertPt);
+  // Old version
+  auto *PtrInGeneric = new AddrSpaceCastInst(PtrInShared, Ptr->getType(),
+                                             Ptr->getName(), &*InsertPt);
+  // Replace with PtrInGeneric all uses of Ptr except PtrInShared.
+  Ptr->replaceAllUsesWith(PtrInGeneric);
+  PtrInShared->setOperand(0, Ptr);
+}
+
+// =============================================================================
+// Main function for this pass.
+// =============================================================================
+bool NVPTXFunctionDataSharing::runOnKernelFunction(Function &F) {
+  bool Modified = false;
+
+  // Skip pass if no data sharing is required.
+  if (!F.hasFnAttribute("has-nvptx-shared-depot"))
+    return false;
+
+  for (auto &B : F) {
+    for (auto &I : B) {
+      auto *AI = dyn_cast<AllocaInst>(&I);
+      if (!AI)
+        continue;
+      if (AI->getType()->isPointerTy() && ptrIsStored(AI)) {
+        markPointerAsShared(AI);
+        Modified = true;
+      }
+    }
+  }
+
+  return Modified;
+}
+
+// Device functions only need to copy byval args into local memory.
+bool NVPTXFunctionDataSharing::runOnDeviceFunction(Function &F) {
+  return true;
+}
+
+bool NVPTXFunctionDataSharing::runOnFunction(Function &F) {
+  return isKernelFunction(F) ? runOnKernelFunction(F) : runOnDeviceFunction(F);
+}
+
+FunctionPass *
+llvm::createNVPTXFunctionDataSharingPass(const NVPTXTargetMachine *TM) {
+  return new NVPTXFunctionDataSharing(TM);
+}

lib/Target/NVPTX/NVPTXInstrInfo.td

def MOV_ADDR64 : NVPTXInst<(outs Int64Regs:$dst), (ins imem:$a),
                            [(set Int64Regs:$dst, (Wrapper tglobaladdr:$a))]>;
 
 // Get pointer to local stack.
 let hasSideEffects = 0 in {
   def MOV_DEPOT_ADDR :    NVPTXInst<(outs Int32Regs:$d), (ins i32imm:$num),
                                      "mov.u32 \t$d, __local_depot$num;", []>;
   def MOV_DEPOT_ADDR_64 : NVPTXInst<(outs Int64Regs:$d), (ins i32imm:$num),
                                     "mov.u64 \t$d, __local_depot$num;", []>;
+  def MOV_SHARED_DEPOT_ADDR :    NVPTXInst<(outs Int32Regs:$d), (ins i32imm:$num),
+                                           "mov.u32 \t$d, __shared_depot$num;", []>;
+  def MOV_SHARED_DEPOT_ADDR_64 : NVPTXInst<(outs Int64Regs:$d), (ins i32imm:$num),
+                                           "mov.u64 \t$d, __shared_depot$num;", []>;
 }
 
 
 // copyPhysreg is hard-coded in NVPTXInstrInfo.cpp
 let IsSimpleMove=1, hasSideEffects=0 in {
   def IMOV1rr :  NVPTXInst<(outs Int1Regs:$dst), (ins Int1Regs:$sss),
                            "mov.pred \t$dst, $sss;", []>;
   def IMOV16rr : NVPTXInst<(outs Int16Regs:$dst), (ins Int16Regs:$sss),

lib/Target/NVPTX/NVPTXLowerAlloca.cpp

 //   %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 ; emits st.local.u32
 //
 // And we will rely on NVPTXInferAddressSpaces to combine the last two
 // instructions.
 //
+// In the case of OpenMP shared variables, perform the same transformation as
+// for local variables but using the shared address space.
+//
 //===----------------------------------------------------------------------===//
 
 #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"
if (skipBasicBlock(BB))
     return false;
 
   bool Changed = false;
   for (auto &I : BB) {
     if (auto allocaInst = dyn_cast<AllocaInst>(&I)) {
       Changed = true;
       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);
+
+      // In the CUDA case, this is always a local address.
+      // In offloading to a device using OpenMP this may be an
+      // address allocated in the shared memory of the device.
+      auto *AddrTy = PointerType::get(ETy, ADDRESS_SPACE_LOCAL);
+      bool PtrIsStored = ptrIsStored(allocaInst);
+      bool RequiresSharedMemory =
+          BB.getParent()->hasFnAttribute("has-nvptx-shared-depot");
+
+      // Handle shared args: currently shared args are declared as
+      // an alloca in LLVM-IR code generation and lowered to
+      // shared memory.
+      if (PtrIsStored && RequiresSharedMemory)
+       AddrTy = PointerType::get(ETy, ADDRESS_SPACE_SHARED);
+
+      auto NewASCToLocal = new AddrSpaceCastInst(allocaInst, AddrTy, "");
+      auto *GenericAddrTy = PointerType::get(ETy, ADDRESS_SPACE_GENERIC);
       auto NewASCToGeneric = new AddrSpaceCastInst(NewASCToLocal,
                                                     GenericAddrTy, "");
       NewASCToLocal->insertAfter(allocaInst);
       NewASCToGeneric->insertAfter(NewASCToLocal);
+
+      // If a value is shared then the additional conversions are required for
+      // correctness.
+      if (PtrIsStored && RequiresSharedMemory) {
+        allocaInst->replaceAllUsesWith(NewASCToGeneric);
+        NewASCToLocal->setOperand(0, allocaInst);
+        continue;
+      }
+
       for (Value::use_iterator UI = allocaInst->use_begin(),
                                 UE = allocaInst->use_end();
             UI != UE; ) {
         // Check Load, Store, GEP, and BitCast Uses on alloca and make them
         // use the converted generic address, in order to expose non-generic
         // addrspacecast to NVPTXInferAddressSpaces. For other types
         // of instructions this is unnecessary and may introduce redundant
         // address cast.
         const auto &AllocaUse = *UI++;
         auto LI = dyn_cast<LoadInst>(AllocaUse.getUser());
         if (LI && LI->getPointerOperand() == allocaInst && !LI->isVolatile()) {
           LI->setOperand(LI->getPointerOperandIndex(), NewASCToGeneric);
           continue;
         }
         auto SI = dyn_cast<StoreInst>(AllocaUse.getUser());
-        if (SI && SI->getPointerOperand() == allocaInst && !SI->isVolatile()) {
-          SI->setOperand(SI->getPointerOperandIndex(), NewASCToGeneric);
+        if (SI && !SI->isVolatile()){
+          unsigned Idx;
+          if (SI->getPointerOperand() == allocaInst)
+            Idx = SI->getPointerOperandIndex();
+          else if (SI->getValueOperand() == allocaInst)
+            Idx = 0;
+          else
             continue;
+          SI->setOperand(Idx, NewASCToGeneric);
         }
         auto GI = dyn_cast<GetElementPtrInst>(AllocaUse.getUser());
         if (GI && GI->getPointerOperand() == allocaInst) {
           GI->setOperand(GI->getPointerOperandIndex(), NewASCToGeneric);
           continue;
         }
         auto BI = dyn_cast<BitCastInst>(AllocaUse.getUser());
         if (BI && BI->getOperand(0) == allocaInst) {

lib/Target/NVPTX/NVPTXLowerSharedFrameIndicesPass.cpp

+//===-- NVPTXLowerSharedFrameIndicesPass.cpp - NVPTX lowering  ------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass lowers the frame indices that need to reference the shared memry
+// depot to the shared framed virtual register VRShared. This pass only performs
+// the lowering if the function has the has-nvptx-shared-depot attribute.
+//
+//===----------------------------------------------------------------------===//
+
+#include "NVPTX.h"
+#include "NVPTXUtilities.h"
+#include "NVPTXRegisterInfo.h"
+#include "NVPTXSubtarget.h"
+#include "NVPTXTargetMachine.h"
+#include "llvm/CodeGen/MachineFrameInfo.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineInstrBuilder.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/TargetFrameLowering.h"
+#include "llvm/CodeGen/TargetRegisterInfo.h"
+#include "llvm/CodeGen/TargetSubtargetInfo.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/MC/MachineLocation.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "nvptx-lower-shared-frame-indices"
+
+namespace {
+class NVPTXLowerSharedFrameIndicesPass : public MachineFunctionPass {
+public:
+  static char ID;
+  NVPTXLowerSharedFrameIndicesPass() : MachineFunctionPass(ID) {}
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+
+private:
+  void calculateSharedFrameObjectOffsets(MachineFunction &Fn);
+};
+}
+
+MachineFunctionPass *llvm::createNVPTXLowerSharedFrameIndicesPass() {
+  return new NVPTXLowerSharedFrameIndicesPass();
+}
+
+char NVPTXLowerSharedFrameIndicesPass::ID = 0;
+
+static bool isSharedFrame(
+      MachineBasicBlock::iterator II,
+      MachineFunction &MF) {
+  MachineInstr &currentMI = *II;
+
+  if (!currentMI.getOperand(0).isReg())
+    return false;;
+
+  bool useSharedFrame = false;
+  unsigned AllocRegisterNumber = currentMI.getOperand(0).getReg();
+
+  for (MachineBasicBlock &MBB : MF) {
+    for (MachineInstr &MI : MBB) {
+      if (MI.getOpcode() == NVPTX::cvta_to_shared_yes_64 ||
+          MI.getOpcode() == NVPTX::cvta_to_shared_yes) {
+        if (AllocRegisterNumber == MI.getOperand(1).getReg()) {
+          useSharedFrame = true;
+          break;
+        }
+      }
+    }
+  }
+  return useSharedFrame;
+}
+
+bool NVPTXLowerSharedFrameIndicesPass::runOnMachineFunction(MachineFunction &MF) {
+  bool Modified = false;
+  bool IsKernel = isKernelFunction(MF.getFunction());
+
+  // Skip pass if function is not the kernel.
+  if (!IsKernel)
+    return Modified;
+
+  // Skip pass if no data sharing is required.
+  // TODO: At the moment there is no limit enforced on the amount
+  // of shared memory used for this scheme. A limit will be
+  // introduced once a fallback scheme that uses global memory is
+  // in place: when the size of the shared depot exceeds a certain
+  // threshhold then the global memory scheme should be used instead.
+  if (!MF.getFunction().hasFnAttribute("has-nvptx-shared-depot"))
+    return Modified;
+
+  SmallVector<int, 16> SharedFrameIndices;
+
+  calculateSharedFrameObjectOffsets(MF);
+
+  for (MachineBasicBlock &MBB : MF) {
+    for (MachineInstr &MI : MBB) {
+      for (unsigned i = 0, e = MI.getNumOperands(); i != e; ++i) {
+        if (!MI.getOperand(i).isFI())
+          continue;
+
+        if (i + 1 >= MI.getNumOperands())
+          continue;
+
+        bool IsSharedFrame = false;
+        int FrameIndex = MI.getOperand(i).getIndex();
+
+        for(int SFI : SharedFrameIndices)
+          if (FrameIndex == SFI)
+            IsSharedFrame = true;
+
+        if (!IsSharedFrame && isSharedFrame(MI, MF)) {
+          SharedFrameIndices.push_back(FrameIndex);
+          IsSharedFrame = true;
+        }
+
+        if (IsSharedFrame) {
+          // Change Frame index to use shared stack.
+          MachineFunction &MF = *MI.getParent()->getParent();
+          int Offset = MF.getFrameInfo().getObjectOffset(FrameIndex) +
+                       MI.getOperand(i + 1).getImm();
+
+          // Using I0 as the frame pointer
+          // For shared data use the appropriate virtual register: VRShared
+          MI.getOperand(i).ChangeToRegister(NVPTX::VRShared, false);
+          MI.getOperand(i + 1).ChangeToImmediate(Offset);
+        }
+        Modified = true;
+      }
+    }
+  }
+
+  return Modified;
+}
+
+/// AdjustStackOffset - Helper function used to adjust the stack frame offset.
+static inline void
+AdjustStackOffset(MachineFrameInfo &MFI, int FrameIdx,
+                  bool StackGrowsDown, int64_t &Offset,
+                  unsigned &MaxAlign) {
+  // If the stack grows down, add the object size to find the lowest address.
+  if (StackGrowsDown)
+    Offset += MFI.getObjectSize(FrameIdx);
+
+  unsigned Align = MFI.getObjectAlignment(FrameIdx);
+
+  // If the alignment of this object is greater than that of the stack, then
+  // increase the stack alignment to match.
+  MaxAlign = std::max(MaxAlign, Align);
+
+  // Adjust to alignment boundary.
+  Offset = (Offset + Align - 1) / Align * Align;
+
+  if (StackGrowsDown) {
+    DEBUG(dbgs() << "alloc FI(" << FrameIdx << ") at SP[" << -Offset << "]\n");
+    MFI.setObjectOffset(FrameIdx, -Offset); // Set the computed offset
+  } else {
+    DEBUG(dbgs() << "alloc FI(" << FrameIdx << ") at SP[" << Offset << "]\n");
+    MFI.setObjectOffset(FrameIdx, Offset);
+    Offset += MFI.getObjectSize(FrameIdx);
+  }
+}
+
+/// This function computes the offset inside the shared stack.
+///
+/// TODO: For simplicity, currently, the offsets conincide with
+/// the local stack frame offsets - the local and stack frame
+/// offsets are the same length.
+void
+NVPTXLowerSharedFrameIndicesPass::calculateSharedFrameObjectOffsets(
+      MachineFunction &Fn) {
+  const TargetFrameLowering &TFI = *Fn.getSubtarget().getFrameLowering();
+  const TargetRegisterInfo *RegInfo = Fn.getSubtarget().getRegisterInfo();
+
+  bool StackGrowsDown =
+    TFI.getStackGrowthDirection() == TargetFrameLowering::StackGrowsDown;
+
+  // Loop over all of the stack objects, assigning sequential addresses...
+  MachineFrameInfo &MFI = Fn.getFrameInfo();
+
+  // Start at the beginning of the local area.
+  // The Offset is the distance from the stack top in the direction
+  // of stack growth -- so it's always nonnegative.
+  int LocalAreaOffset = TFI.getOffsetOfLocalArea();
+  if (StackGrowsDown)
+    LocalAreaOffset = -LocalAreaOffset;
+  assert(LocalAreaOffset >= 0
+         && "Local area offset should be in direction of stack growth");
+  int64_t Offset = LocalAreaOffset;
+
+  // If there are fixed sized objects that are preallocated in the local area,
+  // non-fixed objects can't be allocated right at the start of local area.
+  // We currently don't support filling in holes in between fixed sized
+  // objects, so we adjust 'Offset' to point to the end of last fixed sized
+  // preallocated object.
+  for (int i = MFI.getObjectIndexBegin(); i != 0; ++i) {
+    int64_t FixedOff;
+    if (StackGrowsDown) {
+      // The maximum distance from the stack pointer is at lower address of
+      // the object -- which is given by offset. For down growing stack
+      // the offset is negative, so we negate the offset to get the distance.
+      FixedOff = -MFI.getObjectOffset(i);
+    } else {
+      // The maximum distance from the start pointer is at the upper
+      // address of the object.
+      FixedOff = MFI.getObjectOffset(i) + MFI.getObjectSize(i);
+    }
+    if (FixedOff > Offset) Offset = FixedOff;
+  }
+
+  // NOTE: We do not have a call stack
+
+  unsigned MaxAlign = MFI.getMaxAlignment();
+
+  // No scavenger
+
+  // FIXME: Once this is working, then enable flag will change to a target
+  // check for whether the frame is large enough to want to use virtual
+  // frame index registers. Functions which don't want/need this optimization
+  // will continue to use the existing code path.
+  if (MFI.getUseLocalStackAllocationBlock()) {
+    unsigned Align = MFI.getLocalFrameMaxAlign();
+
+    // Adjust to alignment boundary.
+    Offset = (Offset + Align - 1) / Align * Align;
+
+    DEBUG(dbgs() << "Local frame base offset: " << Offset << "\n");
+
+    // Resolve offsets for objects in the local block.
+    for (unsigned i = 0, e = MFI.getLocalFrameObjectCount(); i != e; ++i) {
+      std::pair<int, int64_t> Entry = MFI.getLocalFrameObjectMap(i);
+      int64_t FIOffset = (StackGrowsDown ? -Offset : Offset) + Entry.second;
+      DEBUG(dbgs() << "alloc FI(" << Entry.first << ") at SP[" <<
+            FIOffset << "]\n");
+      MFI.setObjectOffset(Entry.first, FIOffset);
+    }
+    // Allocate the local block
+    Offset += MFI.getLocalFrameSize();
+
+    MaxAlign = std::max(Align, MaxAlign);
+  }
+
+  // No stack protector
+
+  // Then assign frame offsets to stack objects that are not used to spill
+  // callee saved registers.
+  for (unsigned i = 0, e = MFI.getObjectIndexEnd(); i != e; ++i) {
+    if (MFI.isObjectPreAllocated(i) &&
+        MFI.getUseLocalStackAllocationBlock())
+      continue;
+    if (MFI.isDeadObjectIndex(i))
+      continue;
+
+    AdjustStackOffset(MFI, i, StackGrowsDown, Offset, MaxAlign);
+  }
+
+  // No scavenger
+
+  if (!TFI.targetHandlesStackFrameRounding()) {
+    // If we have reserved argument space for call sites in the function
+    // immediately on entry to the current function, count it as part of the
+    // overall stack size.
+    if (MFI.adjustsStack() && TFI.hasReservedCallFrame(Fn))
+      Offset += MFI.getMaxCallFrameSize();
+
+    // Round up the size to a multiple of the alignment.  If the function has
+    // any calls or alloca's, align to the target's StackAlignment value to
+    // ensure that the callee's frame or the alloca data is suitably aligned;
+    // otherwise, for leaf functions, align to the TransientStackAlignment
+    // value.
+    unsigned StackAlign;
+    if (MFI.adjustsStack() || MFI.hasVarSizedObjects() ||
+        (RegInfo->needsStackRealignment(Fn) && MFI.getObjectIndexEnd() != 0))
+      StackAlign = TFI.getStackAlignment();
+    else
+      StackAlign = TFI.getTransientStackAlignment();
+
+    // If the frame pointer is eliminated, all frame offsets will be relative to
+    // SP not FP. Align to MaxAlign so this works.
+    StackAlign = std::max(StackAlign, MaxAlign);
+    unsigned AlignMask = StackAlign - 1;
+    Offset = (Offset + AlignMask) & ~uint64_t(AlignMask);
+  }
+
+  // Update frame info to pretend that this is part of the stack...
+  int64_t StackSize = Offset - LocalAreaOffset;
+  MFI.setStackSize(StackSize);
+}

lib/Target/NVPTX/NVPTXRegisterInfo.h

public:
   BitVector getReservedRegs(const MachineFunction &MF) const override;
 
   void eliminateFrameIndex(MachineBasicBlock::iterator MI, int SPAdj,
                            unsigned FIOperandNum,
                            RegScavenger *RS = nullptr) const override;
 
   unsigned getFrameRegister(const MachineFunction &MF) const override;
 
+  unsigned getSharedFrameRegister(const MachineFunction &MF) const;
+
   ManagedStringPool *getStrPool() const {
     return const_cast<ManagedStringPool *>(&ManagedStrPool);
   }
 
   const char *getName(unsigned RegNo) const {
     std::stringstream O;
     O << "reg" << RegNo;
     return getStrPool()->getManagedString(O.str().c_str())->c_str();

lib/Target/NVPTX/NVPTXRegisterInfo.cpp

void NVPTXRegisterInfo::eliminateFrameIndex(MachineBasicBlock::iterator II,
   // Using I0 as the frame pointer
   MI.getOperand(FIOperandNum).ChangeToRegister(NVPTX::VRFrame, false);
   MI.getOperand(FIOperandNum + 1).ChangeToImmediate(Offset);
 }
 
 unsigned NVPTXRegisterInfo::getFrameRegister(const MachineFunction &MF) const {
   return NVPTX::VRFrame;
 }
+
+unsigned NVPTXRegisterInfo::getSharedFrameRegister(
+    const MachineFunction &MF) const {
+  return NVPTX::VRShared;
+}

lib/Target/NVPTX/NVPTXRegisterInfo.td

 
 //===----------------------------------------------------------------------===//
 //  Registers
 //===----------------------------------------------------------------------===//
 
 // Special Registers used as stack pointer
 def VRFrame         : NVPTXReg<"%SP">;
 def VRFrameLocal    : NVPTXReg<"%SPL">;
+def VRShared        : NVPTXReg<"%SPS">;
+def VRFrameShared   : NVPTXReg<"%SPSH">;
 
 // Special Registers used as the stack
 def VRDepot  : NVPTXReg<"%Depot">;
+def VRSharedDepot  : NVPTXReg<"%SharedDepot">;
 
 // We use virtual registers, but define a few physical registers here to keep
 // SDAG and the MachineInstr layers happy.
 foreach i = 0-4 in {
   def P#i  : NVPTXReg<"%p"#i>;  // Predicate
   def RS#i : NVPTXReg<"%rs"#i>; // 16-bit
   def R#i  : NVPTXReg<"%r"#i>;  // 32-bit
   def RL#i : NVPTXReg<"%rd"#i>; // 64-bit

 def Float32Regs : NVPTXRegClass<[f32], 32, (add (sequence "F%u", 0, 4))>;
 def Float64Regs : NVPTXRegClass<[f64], 64, (add (sequence "FL%u", 0, 4))>;
 def Int32ArgRegs : NVPTXRegClass<[i32], 32, (add (sequence "ia%u", 0, 4))>;
 def Int64ArgRegs : NVPTXRegClass<[i64], 64, (add (sequence "la%u", 0, 4))>;
 def Float32ArgRegs : NVPTXRegClass<[f32], 32, (add (sequence "fa%u", 0, 4))>;
 def Float64ArgRegs : NVPTXRegClass<[f64], 64, (add (sequence "da%u", 0, 4))>;
 
 // Read NVPTXRegisterInfo.cpp to see how VRFrame and VRDepot are used.
 def SpecialRegs : NVPTXRegClass<[i32], 32, (add VRFrame, VRFrameLocal, VRDepot,
-                                            (sequence "ENVREG%u", 0, 31))>;
+    VRShared, VRFrameShared, VRSharedDepot, (sequence "ENVREG%u", 0, 31))>;

lib/Target/NVPTX/NVPTXTargetMachine.cpp

 void initializeNVVMIntrRangePass(PassRegistry&);
 void initializeNVVMReflectPass(PassRegistry&);
 void initializeGenericToNVVMPass(PassRegistry&);
 void initializeNVPTXAllocaHoistingPass(PassRegistry &);
 void initializeNVPTXAssignValidGlobalNamesPass(PassRegistry&);
 void initializeNVPTXLowerAggrCopiesPass(PassRegistry &);
 void initializeNVPTXLowerArgsPass(PassRegistry &);
 void initializeNVPTXLowerAllocaPass(PassRegistry &);
+void initializeNVPTXFunctionDataSharingPass(PassRegistry &);
 
 } // end namespace llvm
 
 extern "C" void LLVMInitializeNVPTXTarget() {
   // Register the target.
   RegisterTargetMachine<NVPTXTargetMachine32> X(getTheNVPTXTarget32());
   RegisterTargetMachine<NVPTXTargetMachine64> Y(getTheNVPTXTarget64());
 
   // FIXME: This pass is really intended to be invoked during IR optimization,
   // but it's very NVPTX-specific.
   PassRegistry &PR = *PassRegistry::getPassRegistry();
   initializeNVVMReflectPass(PR);
   initializeNVVMIntrRangePass(PR);
   initializeGenericToNVVMPass(PR);
   initializeNVPTXAllocaHoistingPass(PR);
   initializeNVPTXAssignValidGlobalNamesPass(PR);
   initializeNVPTXLowerArgsPass(PR);
   initializeNVPTXLowerAllocaPass(PR);
+  initializeNVPTXFunctionDataSharingPass(PR);
   initializeNVPTXLowerAggrCopiesPass(PR);
 }
 
 static std::string computeDataLayout(bool is64Bit) {
   std::string Ret = "e";
 
   if (!is64Bit)
     Ret += "-p:32:32";
public:
   NVPTXTargetMachine &getNVPTXTargetMachine() const {
     return getTM<NVPTXTargetMachine>();
   }
 
   void addIRPasses() override;
   bool addInstSelector() override;
   void addPostRegAlloc() override;
   void addMachineSSAOptimization() override;
+  void addMachineSSALowering() override;
 
   FunctionPass *createTargetRegisterAllocator(bool) override;
   void addFastRegAlloc(FunctionPass *RegAllocPass) override;
   void addOptimizedRegAlloc(FunctionPass *RegAllocPass) override;
 
 private:
   // If the opt level is aggressive, add GVN; otherwise, add EarlyCSE. This
   // function is only called in opt mode.
if (getOptLevel() != CodeGenOpt::None)
     addPass(createNVPTXImageOptimizerPass());
   addPass(createNVPTXAssignValidGlobalNamesPass());
   addPass(createGenericToNVVMPass());
 
   // NVPTXLowerArgs is required for correctness and should be run right
   // before the address space inference passes.
   addPass(createNVPTXLowerArgsPass(&getNVPTXTargetMachine()));
   if (getOptLevel() != CodeGenOpt::None) {
+    // Add address space inference passes
     addAddressSpaceInferencePasses();
     if (!DisableLoadStoreVectorizer)
       addPass(createLoadStoreVectorizerPass());
     addStraightLineScalarOptimizationPasses();
+  } else {
+    // When the shared depot is generated, even when no optimizations are
+    // used, we need to lower certain alloca instructions to the appropriate
+    // memory type for correctness.
+    addPass(createNVPTXFunctionDataSharingPass(&getNVPTXTargetMachine()));
   }
 
   // === LSR and other generic IR passes ===
   TargetPassConfig::addIRPasses();
   // EarlyCSE is not always strong enough to clean up what LSR produces. For
   // example, GVN can combine
   //
   //   %0 = add %a, %b
void NVPTXPassConfig::addOptimizedRegAlloc(FunctionPass *RegAllocPass) {
   addPass(&StackSlotColoringID);
 
   // FIXME: Needs physical registers
   //addPass(&PostRAMachineLICMID);
 
   printAndVerify("After StackSlotColoring");
 }
 
+void NVPTXPassConfig::addMachineSSALowering() {
+  // Lower shared frame indices.
+  addPass(createNVPTXLowerSharedFrameIndicesPass(), false);
+}
+
 void NVPTXPassConfig::addMachineSSAOptimization() {
   // Pre-ra tail duplication.
   if (addPass(&EarlyTailDuplicateID))
     printAndVerify("After Pre-RegAlloc TailDuplicate");
 
   // Optimize PHIs before DCE: removing dead PHI cycles may make more
   // instructions dead.
   addPass(&OptimizePHIsID);
 
+  // To avoid SSA optimizations on the local frame indices from treating
+  // shared and local frame indices the same, we will lower shared frame
+  // before the optimizations are applied.
+  addMachineSSALowering();
+
   // This pass merges large allocas. StackSlotColoring is a different pass
   // which merges spill slots.
   addPass(&StackColoringID);
 
   // If the target requests it, assign local variables to stack slots relative
   // to one another and simplify frame index references where possible.
   addPass(&LocalStackSlotAllocationID);
 

lib/Target/NVPTX/NVPTXUtilities.h

 //===-- NVPTXUtilities - Utilities -----------------------------*- C++ -*-====//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the declaration of the NVVM specific utility functions.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_TARGET_NVPTX_NVPTXUTILITIES_H
 #define LLVM_LIB_TARGET_NVPTX_NVPTXUTILITIES_H
 
+#include "NVPTXTargetMachine.h"
+#include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Value.h"
 #include <cstdarg>
 #include <set>
 #include <string>
 #include <vector>

 
 bool getMinCTASm(const Function &, unsigned &);
 bool getMaxNReg(const Function &, unsigned &);
 bool isKernelFunction(const Function &);
 
 bool getAlign(const Function &, unsigned index, unsigned &);
 bool getAlign(const CallInst &, unsigned index, unsigned &);
 
+bool ptrIsStored(Value *Ptr);
+
 }
 
 #endif

lib/Target/NVPTX/NVPTXUtilities.cpp

 //===- NVPTXUtilities.cpp - Utility Functions -----------------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains miscellaneous utility functions
 //===----------------------------------------------------------------------===//
 
 #include "NVPTXUtilities.h"
 #include "NVPTX.h"
+#include "llvm/Analysis/CaptureTracking.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/MutexGuard.h"
 #include <algorithm>
 #include <cstring>
 #include <map>
 #include <string>
 #include <vector>
 
 namespace llvm {
 
+#define DEBUG_TYPE "nvptx-utilities"
+
 namespace {
 typedef std::map<std::string, std::vector<unsigned> > key_val_pair_t;
 typedef std::map<const GlobalValue *, key_val_pair_t> global_val_annot_t;
 typedef std::map<const Module *, global_val_annot_t> per_module_annot_t;
 } // anonymous namespace
 
 static ManagedStatic<per_module_annot_t> annotationCache;
 static sys::Mutex Lock;
for (int i = 0, n = alignNode->getNumOperands(); i < n; i++) {
           return false;
         }
       }
     }
   }
   return false;
 }
 
+/// Returns true if there are any instructions storing
+/// the address of this pointer.
+bool ptrIsStored(Value *Ptr) {
+  return PointerMayBeCaptured(Ptr, false, true);
+}
+
 } // namespace llvm

test/CodeGen/NVPTX/insert-shared-depot.ll

+; Check that the shared depot is emitted when the "has-nvptx-shared-depot" is used.
+
+; RUN: llc < %s -march=nvptx -mcpu=sm_20 | FileCheck %s --check-prefixes=PTX32,CHECK
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 | FileCheck %s --check-prefixes=PTX64,CHECK
+
+; CHECK: {{.*}}kernel()
+
+; PTX32: .local .align 8{{.*}}.b8{{.*}}__local_depot0
+; PTX64: .local .align 8{{.*}}.b8{{.*}}__local_depot0
+
+; PTX32: .shared .align 8{{.*}}.b8{{.*}}__shared_depot0
+; PTX64: .shared .align 8{{.*}}.b8{{.*}}__shared_depot0
+
+; PTX32: .reg .b32{{.*}}%SPS;
+; PTX64: .reg .b64{{.*}}%SPS;
+
+; PTX32: .reg .b32{{.*}}%SPSH;
+; PTX64: .reg .b64{{.*}}%SPSH;
+
+; PTX32: mov.u32{{.*}}%SPSH, __shared_depot0;
+; PTX64: mov.u64{{.*}}%SPSH, __shared_depot0;
+
+; PTX32: cvta.shared.u32{{.*}}%SPS, %SPSH;
+; PTX64: cvta.shared.u64{{.*}}%SPS, %SPSH;
+
+target datalayout = "e-i64:64-i128:128-v16:16-v32:32-n16:32:64"
+target triple = "nvptx64-unknown-unknown"
+
+define void @kernel() #0 {
+  %A = alloca i32, align 4
+  %shared_args = alloca i8**, align 8
+  call void @callee(i8*** %shared_args)
+  store i32 10, i32* %A
+  ret void
+}
+
+declare void @callee(i8***)
+
+attributes #0 = {"has-nvptx-shared-depot"}
+
+!nvvm.annotations = !{!0}
+!0 = !{void ()* @kernel, !"kernel", i32 1}

test/CodeGen/NVPTX/lower-alloca-shared.ll

+; Test case for NVPTXLowerAlloca pass.
+
+; RUN: opt < %s -S -nvptx-lower-alloca -infer-address-spaces | 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() #0 {
+; PTX-LABEL: .visible .entry kernel(
+  %A = alloca i32
+; CHECK: addrspacecast i32* %A to i32 addrspace(3)*
+; CHECK: addrspacecast i32 addrspace(3)* %1 to i32*
+; CHECK: store i32 0, i32 addrspace(3)* {{%.+}}
+; PTX: add.u64 {{%rd[0-9]+}}, %SPS, 0;
+; PTX: cvta.to.shared.u64 {{%rd[0-9]+}}, {{%rd[0-9]+}};
+; PTX: st.shared.u32 [{{%rd[0-9]+}}], {{%r[0-9]+}}
+  %shared_args = alloca i32**
+  call void @callee(i32*** %shared_args)
+  %1 = load i32**, i32*** %shared_args
+  %2 = getelementptr inbounds i32*, i32** %1, i64 0
+  store i32* %A, i32** %2
+  store i32 0, i32* %A
+  ret void
+}
+
+declare void @callee(i32***)
+
+attributes #0 = {"has-nvptx-shared-depot"}
+
+!nvvm.annotations = !{!0}
+!0 = !{void ()* @kernel, !"kernel", i32 1}

test/CodeGen/NVPTX/no-shared-depot.ll

+; Test case for when the shared depot does not need to be emitted.
+; The shared depot is not emitted if the function does not have the
+; "has-nvptx-shared-depot" attribute.
+
+; RUN: llc < %s -march=nvptx -mcpu=sm_20 | FileCheck %s --check-prefix=PTX32
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 | FileCheck %s --check-prefix=PTX64
+
+; PTX32: {{.*}}kernel()
+; PTX64: {{.*}}kernel()
+
+; PTX32: .local .align 8{{.*}}.b8{{.*}}__local_depot0
+; PTX64: .local .align 8{{.*}}.b8{{.*}}__local_depot0
+
+; PTX32-NOT: .shared .align 8{{.*}}.b8{{.*}}__shared_depot0
+; PTX64-NOT: .shared .align 8{{.*}}.b8{{.*}}__shared_depot0
+
+; PTX32-NOT: .reg .b32{{.*}}%SPS;
+; PTX64-NOT: .reg .b64{{.*}}%SPS;
+
+; PTX32-NOT: .reg .b32{{.*}}%SPSH;
+; PTX64-NOT: .reg .b64{{.*}}%SPSH;
+
+; PTX32-NOT: mov.u32{{.*}}%SPSH, __shared_depot0;
+; PTX64-NOT: mov.u64{{.*}}%SPSH, __shared_depot0;
+
+; PTX32-NOT: cvta.shared.u32{{.*}}%SPS, %SPSH;
+; PTX64-NOT: cvta.shared.u64{{.*}}%SPS, %SPSH;
+
+target datalayout = "e-i64:64-i128:128-v16:16-v32:32-n16:32:64"
+target triple = "nvptx64-unknown-unknown"
+
+define void @kernel() {
+  %A = alloca i32, align 4
+  %shared_args = alloca i8**, align 8
+  call void @callee(i8*** %shared_args)
+  store i32 10, i32* %A
+  ret void
+}
+
+declare void @callee(i8***)
+
+!nvvm.annotations = !{!0}
+!0 = !{void ()* @kernel, !"kernel", i32 1}

test/CodeGen/NVPTX/nvptx-function-data-sharing.ll

+; Test case for NVPTXFunctionDataSharing pass.
+
+; RUN: opt < %s -S -nvptx-function-data-sharing -infer-address-spaces | 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() #0 {
+; PTX-LABEL: .visible .entry kernel(
+  %A = alloca i32
+; CHECK: addrspacecast i32* %A to i32 addrspace(3)*
+; CHECK: addrspacecast i32 addrspace(3)* %A1 to i32*
+; CHECK: store i32 0, i32 addrspace(3)* {{%.+}}
+; PTX: add.u64 {{%rd[0-9]+}}, %SPS, 0;
+; PTX: cvta.to.shared.u64 {{%rd[0-9]+}}, {{%rd[0-9]+}};
+; PTX: st.shared.u32 [{{%rd[0-9]+}}], {{%r[0-9]+}}
+  %shared_args = alloca i32**
+  call void @callee(i32*** %shared_args)
+  %1 = load i32**, i32*** %shared_args
+  %2 = getelementptr inbounds i32*, i32** %1, i64 0
+  store i32* %A, i32** %2
+  store i32 0, i32* %A
+  ret void
+}
+
+declare void @callee(i32***)
+
+attributes #0 = {"has-nvptx-shared-depot"}
+
+!nvvm.annotations = !{!0}
+!0 = !{void ()* @kernel, !"kernel", i32 1}