extern var as kernel inputs

Author: SaeruHikari

src/clangcxx/src/llvm/ASTConsumer.cpp

@@ -1375,7 +1375,7 @@ auto FunctionBuilderBuilder::build(const clang::FunctionDecl *S, bool allowKerne
                 }
 
                 // collect args
-                for (auto param : params) {
+                auto collect_arg = [&] (const VarDecl* param) -> const luisa::compute::RefExpr* {
                     auto Ty = param->getType();
                     if (auto lcType = db->FindOrAddType(Ty, param->getBeginLoc())) {
                         const luisa::compute::RefExpr *local = nullptr;
@@ -1408,9 +1408,20 @@ auto FunctionBuilderBuilder::build(const clang::FunctionDecl *S, bool allowKerne
                             } break;
                         }
                         stack.SetLocal(param, local);
+                        return local;
                     } else {
                         clangcxx_log_error("unfound arg type: {}", Ty.getAsString());
                     }
+                    return nullptr;
+                };
+                luisa::vector<const clang::VarDecl*> input_params(params.begin(), params.end());
+                if (is_kernel) {
+                    for (const auto *var : db->extern_vars) {
+                        collect_arg(var);
+                    }
+                }
+                for (auto param : params) {
+                    collect_arg(param);
                 }
 
                 // ctor initializers
@@ -1468,6 +1479,17 @@ void RecordDeclStmtHandler::run(const MatchFinder::MatchResult &Result) {
     }
 }
 
+void ExternVarHandler::run(const MatchFinder::MatchResult &Result) {
+    if (const auto *S = Result.Nodes.getNodeAs<clang::VarDecl>("VarDecl")) {
+        bool ignore = false;
+        for (auto Anno : S->specific_attrs<clang::AnnotateAttr>())
+            ignore |= isIgnore(Anno);
+        if (!ignore && S->hasExternalStorage()) {
+            db->extern_vars.emplace_back(S);
+        }
+    }
+}
+
 void GlobalVarHandler::run(const MatchFinder::MatchResult &Result) {
     if (const auto *S = Result.Nodes.getNodeAs<clang::VarDecl>("VarDecl")) {
         bool ignore = false;
@@ -1545,6 +1567,13 @@ ASTConsumerBase::ASTConsumerBase() {
                            unless(isExpansionInSystemHeader()))
                            .bind("VarDecl"),
                        &HandlerForGlobalVar);
+
+    HandlerForExternlVar.db = &db;
+    Matcher.addMatcher(varDecl(
+                           unless(isDefinition()),
+                           unless(isExpansionInSystemHeader()))
+                           .bind("VarDecl"),
+                       &HandlerForExternlVar);
 }
 
 ASTConsumerBase::~ASTConsumerBase() {

src/clangcxx/src/llvm/ASTConsumer.h

@@ -72,6 +72,13 @@ struct GlobalVarHandler : public clang::ast_matchers::MatchFinder::MatchCallback
     TypeDatabase *db = nullptr;
 };
 
+struct ExternVarHandler : public clang::ast_matchers::MatchFinder::MatchCallback {
+    ExternVarHandler() = default;
+    void run(const MatchFinder::MatchResult &Result) final;
+
+    TypeDatabase *db = nullptr;
+};
+
 struct FunctionDeclStmtHandler : public clang::ast_matchers::MatchFinder::MatchCallback {
     FunctionDeclStmtHandler() = default;
     void run(const MatchFinder::MatchResult &Result) final;
@@ -89,6 +96,7 @@ class ASTConsumerBase : public clang::ASTConsumer {
     TypeDatabase db;
     RecordDeclStmtHandler HandlerForTypeDecl;
     GlobalVarHandler HandlerForGlobalVar;
+    ExternVarHandler HandlerForExternlVar;
     FunctionDeclStmtHandler HandlerForFuncionDecl;
     clang::ast_matchers::MatchFinder Matcher;
 };

src/clangcxx/src/llvm/TypeDatabase.h

@@ -40,6 +40,7 @@ struct TypeDatabase {
     void DumpWithLocation(const clang::Decl *decl);
 
     // luisa::unordered_map<luisa::string, const luisa::compute::RefExpr *> globals;
+    luisa::vector<const clang::VarDecl*> extern_vars;
     luisa::shared_ptr<compute::detail::FunctionBuilder> kernel_builder;
     luisa::shared_ptr<compute::detail::FunctionBuilder> vertex_builder;
     luisa::shared_ptr<compute::detail::FunctionBuilder> pixel_builder;