CUDA: IR generation support for device stubs

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@141304 91177308-0d34-0410-b5e6-96231b3b80d8

3 files changed, 101 insertions, 0 deletions

diff --git a/lib/CodeGen/CGCUDANV.cpp b/lib/CodeGen/CGCUDANV.cpp
index 69f75a6792..88a0bdc821 100644
--- a/lib/CodeGen/CGCUDANV.cpp
+++ b/lib/CodeGen/CGCUDANV.cpp
@@ -13,6 +13,15 @@
 //===----------------------------------------------------------------------===//
 
 #include "CGCUDARuntime.h"
+#include "CodeGenFunction.h"
+#include "CodeGenModule.h"
+#include "clang/AST/Decl.h"
+#include "llvm/BasicBlock.h"
+#include "llvm/Constants.h"
+#include "llvm/DerivedTypes.h"
+#include "llvm/Support/CallSite.h"
+
+#include <vector>
 
 using namespace clang;
 using namespace CodeGen;
@@ -20,13 +29,96 @@ using namespace CodeGen;
 namespace {
 
 class CGNVCUDARuntime : public CGCUDARuntime {
+
+private:
+  llvm::Type *IntTy, *SizeTy;
+  llvm::PointerType *CharPtrTy, *VoidPtrTy;
+
+  llvm::Constant *getSetupArgumentFn() const;
+  llvm::Constant *getLaunchFn() const;
+
 public:
   CGNVCUDARuntime(CodeGenModule &CGM);
+
+  void EmitDeviceStubBody(CodeGenFunction &CGF, FunctionArgList &Args);
 };
 
 }
 
 CGNVCUDARuntime::CGNVCUDARuntime(CodeGenModule &CGM) : CGCUDARuntime(CGM) {
+  CodeGen::CodeGenTypes &Types = CGM.getTypes();
+  ASTContext &Ctx = CGM.getContext();
+
+  IntTy = Types.ConvertType(Ctx.IntTy);
+  SizeTy = Types.ConvertType(Ctx.getSizeType());
+
+  CharPtrTy = llvm::PointerType::getUnqual(Types.ConvertType(Ctx.CharTy));
+  VoidPtrTy = cast<llvm::PointerType>(Types.ConvertType(Ctx.VoidPtrTy));
+}
+
+llvm::Constant *CGNVCUDARuntime::getSetupArgumentFn() const {
+  // cudaError_t cudaSetupArgument(void *, size_t, size_t)
+  std::vector<llvm::Type*> Params;
+  Params.push_back(VoidPtrTy);
+  Params.push_back(SizeTy);
+  Params.push_back(SizeTy);
+  return CGM.CreateRuntimeFunction(llvm::FunctionType::get(IntTy,
+                                                           Params, false),
+                                   "cudaSetupArgument");
+}
+
+llvm::Constant *CGNVCUDARuntime::getLaunchFn() const {
+  // cudaError_t cudaLaunch(char *)
+  std::vector<llvm::Type*> Params;
+  Params.push_back(CharPtrTy);
+  return CGM.CreateRuntimeFunction(llvm::FunctionType::get(IntTy,
+                                                           Params, false),
+                                   "cudaLaunch");
+}
+
+void CGNVCUDARuntime::EmitDeviceStubBody(CodeGenFunction &CGF,
+                                         FunctionArgList &Args) {
+  // Build the argument value list and the argument stack struct type.
+  llvm::SmallVector<llvm::Value *, 16> ArgValues;
+  std::vector<llvm::Type *> ArgTypes;
+  for (FunctionArgList::const_iterator I = Args.begin(), E = Args.end();
+       I != E; ++I) {
+    llvm::Value *V = CGF.GetAddrOfLocalVar(*I);
+    ArgValues.push_back(V);
+    assert(isa<llvm::PointerType>(V->getType()) && "Arg type not PointerType");
+    ArgTypes.push_back(cast<llvm::PointerType>(V->getType())->getElementType());
+  }
+  llvm::StructType *ArgStackTy = llvm::StructType::get(
+      CGF.getLLVMContext(), ArgTypes);
+
+  llvm::BasicBlock *EndBlock = CGF.createBasicBlock("setup.end");
+
+  // Emit the calls to cudaSetupArgument
+  llvm::Constant *cudaSetupArgFn = getSetupArgumentFn();
+  for (unsigned I = 0, E = Args.size(); I != E; ++I) {
+    llvm::Value *Args[3];
+    llvm::BasicBlock *NextBlock = CGF.createBasicBlock("setup.next");
+    Args[0] = CGF.Builder.CreatePointerCast(ArgValues[I], VoidPtrTy);
+    Args[1] = CGF.Builder.CreateIntCast(
+        llvm::ConstantExpr::getSizeOf(ArgTypes[I]),
+        SizeTy, false);
+    Args[2] = CGF.Builder.CreateIntCast(
+        llvm::ConstantExpr::getOffsetOf(ArgStackTy, I),
+        SizeTy, false);
+    llvm::CallSite CS = CGF.EmitCallOrInvoke(cudaSetupArgFn, Args);
+    llvm::Constant *Zero = llvm::ConstantInt::get(IntTy, 0);
+    llvm::Value *CSZero = CGF.Builder.CreateICmpEQ(CS.getInstruction(), Zero);
+    CGF.Builder.CreateCondBr(CSZero, NextBlock, EndBlock);
+    CGF.EmitBlock(NextBlock);
+  }
+
+  // Emit the call to cudaLaunch
+  llvm::Constant *cudaLaunchFn = getLaunchFn();
+  llvm::Value *Arg = CGF.Builder.CreatePointerCast(CGF.CurFn, CharPtrTy);
+  CGF.EmitCallOrInvoke(cudaLaunchFn, Arg);
+  CGF.EmitBranch(EndBlock);
+
+  CGF.EmitBlock(EndBlock);
 }
 
 CGCUDARuntime *CodeGen::CreateNVCUDARuntime(CodeGenModule &CGM) {
diff --git a/lib/CodeGen/CGCUDARuntime.h b/lib/CodeGen/CGCUDARuntime.h
index 79890e89d6..a99a67ae1a 100644
--- a/lib/CodeGen/CGCUDARuntime.h
+++ b/lib/CodeGen/CGCUDARuntime.h
@@ -24,6 +24,7 @@ namespace CodeGen {
 
 class CodeGenFunction;
 class CodeGenModule;
+class FunctionArgList;
 class ReturnValueSlot;
 class RValue;
 
@@ -39,6 +40,9 @@ public:
                                         const CUDAKernelCallExpr *E,
                                         ReturnValueSlot ReturnValue);
   
+  virtual void EmitDeviceStubBody(CodeGenFunction &CGF,
+                                  FunctionArgList &Args) = 0;
+
 };
 
 /// Creates an instance of a CUDA runtime class.
diff --git a/lib/CodeGen/CodeGenFunction.cpp b/lib/CodeGen/CodeGenFunction.cpp
index 4c6624f761..12def6728f 100644
--- a/lib/CodeGen/CodeGenFunction.cpp
+++ b/lib/CodeGen/CodeGenFunction.cpp
@@ -13,6 +13,7 @@
 
 #include "CodeGenFunction.h"
 #include "CodeGenModule.h"
+#include "CGCUDARuntime.h"
 #include "CGCXXABI.h"
 #include "CGDebugInfo.h"
 #include "CGException.h"
@@ -404,6 +405,10 @@ void CodeGenFunction::GenerateCode(GlobalDecl GD, llvm::Function *Fn,
     EmitDestructorBody(Args);
   else if (isa<CXXConstructorDecl>(FD))
     EmitConstructorBody(Args);
+  else if (getContext().getLangOptions().CUDA &&
+           !CGM.getCodeGenOpts().CUDAIsDevice &&
+           FD->hasAttr<CUDAGlobalAttr>())
+    CGM.getCUDARuntime().EmitDeviceStubBody(*this, Args);
   else
     EmitFunctionBody(Args);