Add support for preheaders in KernelBuilder

Author: stijnh

include/kernel_launcher/kernel.h

@@ -91,6 +91,7 @@ struct KernelBuilder: ConfigSpace {
     KernelBuilder& assertion(TypedExpr<bool> e);
     KernelBuilder& define(std::string name, TypedExpr<std::string> value);
     KernelBuilder& compiler_flag(TypedExpr<std::string> opt);
+    KernelBuilder& include_header(KernelSource source);
 
     template<typename T, typename... Ts>
     KernelBuilder& template_args(T&& first, Ts&&... rest) {
@@ -152,6 +153,7 @@ struct KernelBuilder: ConfigSpace {
   private:
     std::string kernel_name_;
     KernelSource kernel_source_;
+    std::vector<KernelSource> preheaders_;
     std::array<TypedExpr<uint32_t>, 3> block_size_ = {1u, 1u, 1u};
     std::array<TypedExpr<uint32_t>, 3> grid_size_ = {1u, 1u, 1u};
     TypedExpr<uint32_t> shared_mem_ = {0u};

include/kernel_launcher/registry.h

@@ -21,13 +21,14 @@ struct KernelDescriptor {
 
 struct AnyKernelDescriptor {
     AnyKernelDescriptor(AnyKernelDescriptor&&) noexcept = default;
+    AnyKernelDescriptor(AnyKernelDescriptor&) noexcept = default;
     AnyKernelDescriptor(const AnyKernelDescriptor&) = default;
 
     template<typename D>
     AnyKernelDescriptor(D&& descriptor) {
         using T = typename std::decay<D>::type;
         descriptor_type_ = type_of<T>();
-        descriptor_ = std::make_unique<T>(std::forward<D>(descriptor));
+        descriptor_ = std::make_shared<T>(std::forward<D>(descriptor));
         hash_ = hash_fields(descriptor_type_, descriptor_->hash());
     }
 
@@ -53,16 +54,16 @@ struct AnyKernelDescriptor {
     TypeInfo descriptor_type_;
     std::shared_ptr<KernelDescriptor> descriptor_;
 };
-}
+}  // namespace kernel_launcher
 
 namespace std {
-template <>
+template<>
 struct hash<kernel_launcher::AnyKernelDescriptor> {
     size_t operator()(const kernel_launcher::AnyKernelDescriptor& d) const {
         return d.hash();
     }
 };
-}
+}  // namespace std
 
 namespace kernel_launcher {
 struct KernelRegistry {
@@ -108,19 +109,16 @@ struct KernelRegistry {
 
 const KernelRegistry& default_registry();
 
-template<typename D>
-WisdomKernelLaunch
-launch(D&& descriptor, cudaStream_t stream, ProblemSize size) {
-    return default_registry().instantiate(
-        std::forward<D>(descriptor, stream, size));
+inline WisdomKernelLaunch
+launch(AnyKernelDescriptor descriptor, cudaStream_t stream, ProblemSize size) {
+    return default_registry().instantiate(descriptor, stream, size);
 }
 
-template<typename D>
-WisdomKernelLaunch launch(D&& descriptor, ProblemSize size) {
-    return launch(std::forward<D>(descriptor), (cudaStream_t) nullptr, size);
+inline WisdomKernelLaunch
+launch(AnyKernelDescriptor descriptor, ProblemSize size) {
+    return launch(descriptor, (cudaStream_t) nullptr, size);
 }
 
 }  // namespace kernel_launcher
 
-
 #endif  //KERNEL_LAUNCHER_CACHE_H

src/export.cpp

@@ -175,7 +175,20 @@ struct KernelBuilderSerializerHack {
     }
 
     static json builder_to_json(const KernelBuilder& builder) {
-        std::unordered_map<std::string, json> defines;
+        std::vector<json> headers;
+        for (const auto& source : builder.preheaders_) {
+            json content = nullptr;
+            if (source.content() != nullptr) {
+                content = *source.content();
+            }
+
+            headers.push_back({
+                {"file", source.file_name()},
+                {"content", std::move(content)},
+            });
+        }
+
+        json defines;
         for (const auto& p : builder.defines_) {
             defines[p.first] = expr_to_json(p.second);
         }
@@ -195,6 +208,7 @@ struct KernelBuilderSerializerHack {
         result["shared_memory"] = expr_to_json(builder.shared_mem_);
         result["template_args"] = expr_list_to_json(builder.template_args_);
         result["defines"] = std::move(defines);
+        result["headers"] = std::move(headers);
 
         return result;
     }

src/kernel.cpp

@@ -58,6 +58,11 @@ KernelBuilder& KernelBuilder::compiler_flag(TypedExpr<std::string> opt) {
     return *this;
 }
 
+KernelBuilder& KernelBuilder::include_header(KernelSource source) {
+    preheaders_.push_back(std::move(source));
+    return *this;
+}
+
 KernelDef KernelBuilder::build(
     const Config& config,
     const std::vector<TypeInfo>& param_types) const {
@@ -86,6 +91,10 @@ KernelDef KernelBuilder::build(
         }
     }
 
+    for (const auto& source : preheaders_) {
+        def.add_preincluded_header(source);
+    }
+
     for (const auto& p : defines_) {
         def.add_compiler_option("--define-macro");
         def.add_compiler_option(p.first + "=" + eval(p.second));