slang error on kernel launch from another kernel

Bug: 27459075

Extended the AST visitor that lowers rsForEach() and
rsForEachWithOptions() calls to be aware if the function being walked
is a kernel itself. If so, on seeing a call to one of the two kernel
launch functions, reports a compiler error.

Note that such checking is oblivious to the control flow.

Also removed all checks inside the LowerRSForEachCall() function,
since we do all the same checks before entering that function.

Change-Id: Ib7f276e8e8d707f39b310548c084ed550a7a3480

7 files changed, 60 insertions, 23 deletions

diff --git a/slang_backend.cpp b/slang_backend.cpp
index cbfdb32..c6b5a58 100644
--- a/slang_backend.cpp
+++ b/slang_backend.cpp
@@ -371,20 +371,6 @@ void Backend::AnnotateFunction(clang::FunctionDecl *FD) {
   }
 }
 
-void Backend::LowerRSForEachCall(clang::FunctionDecl *FD) {
-  // Skip this AST walking for lower API levels.
-  if (getTargetAPI() < SLANG_DEVELOPMENT_TARGET_API) {
-    return;
-  }
-
-  if (!FD || !FD->hasBody() ||
-      Slang::IsLocInRSHeaderFile(FD->getLocation(), mSourceMgr)) {
-    return;
-  }
-
-  mForEachHandler.VisitStmt(FD->getBody());
-}
-
 bool Backend::HandleTopLevelDecl(clang::DeclGroupRef D) {
   // Find and remember the types for rs_allocation and rs_script_call_t so
   // they can be used later for translating rsForEach() calls.
@@ -446,17 +432,16 @@ bool Backend::HandleTopLevelDecl(clang::DeclGroupRef D) {
 
     if (getTargetAPI() == SLANG_DEVELOPMENT_TARGET_API) {
       if (FD && FD->hasBody() &&
-          RSExportForEach::isRSForEachFunc(getTargetAPI(), FD)) {
-        // Log kernels by their names, and assign them slot numbers.
-        if (!Slang::IsLocInRSHeaderFile(FD->getLocation(), mSourceMgr)) {
-            mContext->addForEach(FD);
+          !Slang::IsLocInRSHeaderFile(FD->getLocation(), mSourceMgr)) {
+        bool isKernel = RSExportForEach::isRSForEachFunc(getTargetAPI(), FD);
+        if (isKernel) {
+          // Log kernels by their names, and assign them slot numbers.
+          mContext->addForEach(FD);
         }
-      } else {
         // Look for any kernel launch calls and translate them into using the
         // internal API.
-        // TODO: Simply ignores kernel launch inside a kernel for now.
-        // Needs more rigorous and comprehensive checks.
-        LowerRSForEachCall(FD);
+        // Report a compiler on kernel launches inside a kernel.
+        mForEachHandler.handleForEachCalls(FD, getTargetAPI());
       }
     }
   }
diff --git a/slang_backend.h b/slang_backend.h
index 3c9b495..ef7f635 100644
--- a/slang_backend.h
+++ b/slang_backend.h
@@ -114,7 +114,12 @@ class Backend : public clang::ASTConsumer {
   void dumpExportReduceNewInfo(llvm::Module *M);
   void dumpExportTypeInfo(llvm::Module *M);
 
-  void LowerRSForEachCall(clang::FunctionDecl* FD);
+  // Translates any rsForEach() or rsForEachWithOptions() calls inside the body
+  // of FD to lower-level runtime calls to rsForEachInternal(), if FD is not a
+  // kernel function itself, as indicated by isKernel being false. If isKernel
+  // is true, reports an error on any calls to rsForEach() or
+  // rsForEachWithOptions().
+  void LowerRSForEachCall(clang::FunctionDecl* FD, bool isKernel);
 
  protected:
   llvm::LLVMContext &mLLVMContext;
diff --git a/slang_rs_foreach_lowering.cpp b/slang_rs_foreach_lowering.cpp
index bbaaa7b..05553a7 100644
--- a/slang_rs_foreach_lowering.cpp
+++ b/slang_rs_foreach_lowering.cpp
@@ -110,6 +110,12 @@ const clang::FunctionDecl* RSForEachLowering::matchKernelLaunchCall(
     return nullptr;
   }
 
+  if (mInsideKernel) {
+    mCtxt->ReportError(CE->getExprLoc(),
+        "Invalid kernel launch call made from inside another kernel.");
+    return nullptr;
+  }
+
   clang::Expr* arg0 = CE->getArg(0);
   const clang::FunctionDecl* kernel = matchFunctionDesignator(arg0);
 
@@ -356,4 +362,12 @@ void RSForEachLowering::VisitStmt(clang::Stmt* S) {
   }
 }
 
+void RSForEachLowering::handleForEachCalls(clang::FunctionDecl* FD,
+                                           unsigned int targetAPI) {
+  slangAssert(FD && FD->hasBody());
+
+  mInsideKernel = RSExportForEach::isRSForEachFunc(targetAPI, FD);
+  VisitStmt(FD->getBody());
+}
+
 }  // namespace slang
diff --git a/slang_rs_foreach_lowering.h b/slang_rs_foreach_lowering.h
index 2e30733..56ec66b 100644
--- a/slang_rs_foreach_lowering.h
+++ b/slang_rs_foreach_lowering.h
@@ -34,12 +34,23 @@ class RSForEachLowering : public clang::StmtVisitor<RSForEachLowering> {
  public:
   RSForEachLowering(RSContext* ctxt);
 
+  // Given a FunctionDecl FD and the target API level, either translates all
+  // rsForEach() and rsForEachWithOptions() calls inside FD into calls to the
+  // low-level rsForEachInternal() API, if FD is not a kernel function itself;
+  // or, in the case where FD is a kernel function, reports a compiler error on
+  // any calls to either kernel launching API function.
+  void handleForEachCalls(clang::FunctionDecl* FD, unsigned int targetAPI);
+
   void VisitCallExpr(clang::CallExpr *CE);
   void VisitStmt(clang::Stmt *S);
 
  private:
   RSContext* mCtxt;
   clang::ASTContext& mASTCtxt;
+  // A flag, if true, indicating that the visitor is walking inside a kernel
+  // function, in which case any call to rsForEach() or rsForEachWithOptions()
+  // is a compiler error.
+  bool mInsideKernel;
 
   const clang::FunctionDecl* matchFunctionDesignator(clang::Expr* expr);
   const clang::FunctionDecl* matchKernelLaunchCall(clang::CallExpr* CE,
diff --git a/tests/F_foreach_from_kernel/foreach_from_kernel.rs b/tests/F_foreach_from_kernel/foreach_from_kernel.rs
new file mode 100644
index 0000000..911170c
--- /dev/null
+++ b/tests/F_foreach_from_kernel/foreach_from_kernel.rs
@@ -0,0 +1,21 @@
+// -target-api 0
+#pragma version(1)
+#pragma rs java_package_name(com.example.foo)
+
+int __attribute__((kernel)) foo(int a) {
+  return a;
+}
+
+rs_allocation g_in;
+rs_allocation g_out;
+
+int __attribute__((kernel)) goo(int a) {
+  rsForEach(foo, g_in, g_out);
+  return a;
+}
+
+void bar() {
+  rsForEach(foo, g_in, g_out);
+}
+
+
diff --git a/tests/F_foreach_from_kernel/stderr.txt.expect b/tests/F_foreach_from_kernel/stderr.txt.expect
new file mode 100644
index 0000000..d9b82b0
--- /dev/null
+++ b/tests/F_foreach_from_kernel/stderr.txt.expect
@@ -0,0 +1 @@
+foreach_from_kernel.rs:13:3: error: Invalid kernel launch call made from inside another kernel.
diff --git a/tests/F_foreach_from_kernel/stdout.txt.expect b/tests/F_foreach_from_kernel/stdout.txt.expect
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/tests/F_foreach_from_kernel/stdout.txt.expect