From 124a88ac597b2a1151ec62b835c53e01bfa56b1f Mon Sep 17 00:00:00 2001
From: Vladimir Lazarev <vladimir.lazarev@intel.com>
Date: Fri, 28 Dec 2018 22:52:26 +0300
Subject: [PATCH] [SYCL] Mark functions that the kernels call with SYCL device
 attribute.

Signed-off-by: Vladimir Lazarev <vladimir.lazarev@intel.com>
---
 clang/lib/CodeGen/CodeGenModule.cpp         |  6 +---
 clang/lib/Sema/SemaSYCL.cpp                 | 34 +++++++++++++++++++++
 clang/test/CodeGenSYCL/device-functions.cpp | 27 ++++++++++++++++
 3 files changed, 62 insertions(+), 5 deletions(-)
 create mode 100644 clang/test/CodeGenSYCL/device-functions.cpp

diff --git a/clang/lib/CodeGen/CodeGenModule.cpp b/clang/lib/CodeGen/CodeGenModule.cpp
index 031ece2ba1cd0..68f1b1278de36 100644
--- a/clang/lib/CodeGen/CodeGenModule.cpp
+++ b/clang/lib/CodeGen/CodeGenModule.cpp
@@ -2878,11 +2878,7 @@ llvm::Constant *CodeGenModule::GetOrCreateLLVMFunction(
       // Look for a declaration that's lexically in a record.
       for (const auto *FD = cast<FunctionDecl>(D)->getMostRecentDecl(); FD;
            FD = FD->getPreviousDecl()) {
-        // For SYCL we also should emit a definition for a function because all
-        // top-level declarations without OpenCL kernel attribute are ignored
-        // now.
-        // TODO: fix this hack
-        if (isa<CXXRecordDecl>(FD->getLexicalDeclContext()) || LangOpts.SYCL) {
+        if (isa<CXXRecordDecl>(FD->getLexicalDeclContext())) {
           if (FD->doesThisDeclarationHaveABody()) {
             addDeferredDeclToEmit(GD.getWithDecl(FD));
             break;
diff --git a/clang/lib/Sema/SemaSYCL.cpp b/clang/lib/Sema/SemaSYCL.cpp
index 8611c7884d852..6ea0d5f6e2827 100644
--- a/clang/lib/Sema/SemaSYCL.cpp
+++ b/clang/lib/Sema/SemaSYCL.cpp
@@ -13,6 +13,7 @@
 #include "clang/Sema/Sema.h"
 #include "llvm/ADT/SmallVector.h"
 #include "TreeTransform.h"
+#include "clang/AST/RecursiveASTVisitor.h"
 
 using namespace clang;
 
@@ -28,6 +29,35 @@ enum target {
   image_array
 };
 
+class MarkDeviceFunction : public RecursiveASTVisitor<MarkDeviceFunction> {
+public:
+  MarkDeviceFunction(Sema &S)
+      : RecursiveASTVisitor<MarkDeviceFunction>(), SemaRef(S) {}
+  bool VisitCallExpr(CallExpr *e) {
+    if (FunctionDecl *Callee = e->getDirectCallee()) {
+      // Remember that all SYCL kernel functions have deferred
+      // instantiation as template functions. It means that
+      // all functions used by kernel have already been parsed and have
+      // definitions.
+      if (FunctionDecl *Def = Callee->getDefinition()) {
+        if (!Def->hasAttr<SYCLDeviceAttr>()) {
+          Def->addAttr(SYCLDeviceAttr::CreateImplicit(SemaRef.Context));
+          this->TraverseStmt(Def->getBody());
+          // But because parser works with top level decls and codegen
+          // already saw and ignored our function without device attribute we
+          // need to add this function into sycl kernels array to show it
+          // this function again.
+          SemaRef.AddSyclKernel(Def);
+        }
+      }
+    }
+    return true;
+  }
+
+private:
+  Sema &SemaRef;
+};
+
 class KernelBodyTransform : public TreeTransform<KernelBodyTransform> {
 public:
   KernelBodyTransform(llvm::DenseMap<DeclaratorDecl *, DeclaratorDecl *> &Map,
@@ -302,6 +332,10 @@ void Sema::ConstructSYCLKernel(FunctionDecl *KernelHelper) {
     SYCLKernel->setBody(SYCLKernelBody);
 
     AddSyclKernel(SYCLKernel);
+
+    // Let's mark all called functions with SYCL Device attribute.
+    MarkDeviceFunction Marker(*this);
+    Marker.TraverseStmt(SYCLKernelBody);
   }
 }
 
diff --git a/clang/test/CodeGenSYCL/device-functions.cpp b/clang/test/CodeGenSYCL/device-functions.cpp
new file mode 100644
index 0000000000000..082572742ec9e
--- /dev/null
+++ b/clang/test/CodeGenSYCL/device-functions.cpp
@@ -0,0 +1,27 @@
+// RUN: %clang -cc1 -triple spir64-unknown-linux-sycldevice -std=c++11 -fsycl-is-device -S -emit-llvm -x c++ %s -o - | FileCheck %s
+
+template <typename T>
+T bar(T arg);
+
+void foo() {
+  int a = 1 + 1 + bar(1);
+}
+
+template <typename T>
+T bar(T arg) {
+  return arg;
+}
+
+template <typename name, typename Func>
+__attribute__((sycl_kernel)) void kernel_single_task(Func kernelFunc) {
+  kernelFunc();
+}
+
+int main() {
+  kernel_single_task<class fake_kernel>([]() { foo(); });
+  return 0;
+}
+// CHECK: define spir_kernel void @fake_kernel()
+// CHECK: define internal spir_func void @"_ZZ4mainENK3$_0clEv"(%class.anon* %this)
+// CHECK: define spir_func void @_Z3foov()
+// CHECK: define linkonce_odr spir_func i32 @_Z3barIiET_S0_(i32 %arg)