RTC: only inlucde the required message types for a given kernel.

This provides a suprising improvment to RTC build time
FLAMEGPU · Jul 1, 2021 · 7a5f81b · 7a5f81b
1 parent ad9a28b
commit 7a5f81b
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/include/flamegpu/runtime/DeviceAPI.cuh b/include/flamegpu/runtime/DeviceAPI.cuh
@@ -8,14 +8,14 @@
 
 #ifndef __CUDACC_RTC__
 #include "flamegpu/runtime/detail/curve/curve.cuh"
+#include "flamegpu/runtime/messaging_device.h"
 #else
 #include "dynamic/curve_rtc_dynamic.h"
 #endif  // !_RTC
 #include "flamegpu/runtime/utility/AgentRandom.cuh"
 #include "flamegpu/runtime/utility/DeviceEnvironment.cuh"
 #include "flamegpu/runtime/AgentFunction.cuh"
 #include "flamegpu/runtime/AgentFunctionCondition.cuh"
-#include "flamegpu/runtime/messaging_device.h"
 #include "flamegpu/defines.h"
 
 namespace flamegpu {

diff --git a/src/flamegpu/model/AgentFunctionDescription.cpp b/src/flamegpu/model/AgentFunctionDescription.cpp
@@ -482,17 +482,26 @@ bool AgentFunctionDescription::isRTC() const {
 
 AgentFunctionDescription& AgentDescription::newRTCFunction(const std::string& function_name, const std::string& func_src) {
     if (agent->functions.find(function_name) == agent->functions.end()) {
-        // append jitify program string and include
-        std::string func_src_str = std::string(function_name + "_program\n").append("#include \"flamegpu/runtime/DeviceAPI.cuh\"\n").append(func_src);
         // Use Regex to get agent function name, and input/output message type
         std::regex rgx(R"###(.*FLAMEGPU_AGENT_FUNCTION\([ \t]*(\w+),[ \t]*([:\w]+),[ \t]*([:\w]+)[ \t]*\))###");
         std::smatch match;
-        if (std::regex_search(func_src_str, match, rgx)) {
+        if (std::regex_search(func_src, match, rgx)) {
             if (match.size() == 4) {
                 std::string code_func_name = match[1];  // not yet clear if this is required
                 std::string in_type_name = match[2];
                 std::string out_type_name = match[3];
                 // set the runtime agent function source in agent function data
+                std::string func_src_str = std::string(function_name + "_program\n").append("#include \"flamegpu/runtime/DeviceAPI.cuh\"\n");
+                // Include the required headers for the input message type.
+                std::string in_type_include_name = in_type_name.substr(in_type_name.find_last_of("Msg") + 1);
+                func_src_str = func_src_str.append("#include \"flamegpu/runtime/messaging/"+ in_type_include_name + "/" + in_type_include_name + "Device.cuh\"\n");
+                // If the message input and output types do not match, also include the input type
+                if(in_type_name != out_type_name) {
+                    std::string out_type_include_name = out_type_name.substr(in_type_name.find_last_of("Msg") + 1);
+                    func_src_str = func_src_str.append("#include \"flamegpu/runtime/messaging/"+ out_type_include_name + "/" + out_type_include_name + "Device.cuh\"\n");
+                }
+                // Append the function source
+                func_src_str = func_src_str.append(func_src);
                 auto rtn = std::shared_ptr<AgentFunctionData>(new AgentFunctionData(this->agent->shared_from_this(), function_name, func_src_str, in_type_name, out_type_name, code_func_name));
                 agent->functions.emplace(function_name, rtn);
                 return *rtn->description;