basic work-graph codegen

Author: MaxwellGengYF

src/backends/common/hlsl/CMakeLists.txt

@@ -56,6 +56,7 @@ if (LUISA_COMPUTE_ENABLE_DX OR LUISA_COMPUTE_ENABLE_VULKAN)
             codegen_utils/constant.cpp
             codegen_utils/entry_points.cpp
             codegen_utils/function_codegen.cpp
+            codegen_utils/work_graph_codegen.cpp
             codegen_utils/property.cpp
             codegen_utils/resource.cpp
             codegen_utils/type_system.cpp

src/backends/common/hlsl/builtin/hlsl_builtin.hpp

@@ -23,6 +23,7 @@
 
 extern "C" {
 LC_HLSL_DECL_VARNAME(hlsl_header_bytes)
+LC_HLSL_DECL_VARNAME(work_graph_bytes)
 LC_HLSL_DECL_VARNAME(dx_linalg_bytes)
 LC_HLSL_DECL_VARNAME(hlsl_header_fallback_bytes)
 LC_HLSL_DECL_VARNAME(raytracing_header_bytes)
@@ -62,6 +63,7 @@ static HLSLCompressedHeader get_hlsl_builtin(luisa::string_view ss) {
         luisa::unordered_map<luisa::string_view, HLSLCompressedHeader> dict;
         Dict() {
             LC_HLSL_INSERT_VARNAME(hlsl_header_bytes, "hlsl_header")
+            LC_HLSL_INSERT_VARNAME(work_graph_bytes, "work_graph")
             LC_HLSL_INSERT_VARNAME(spv_alias_bytes, "spv_alias")
             LC_HLSL_INSERT_VARNAME(dx_linalg_bytes, "dx_linalg")
             LC_HLSL_INSERT_VARNAME(hlsl_header_fallback_bytes, "hlsl_header_fallback")

src/backends/common/hlsl/codegen_utils/entry_points.cpp

@@ -447,9 +447,9 @@ CodegenResult CodegenUtility::WorkGraphCodegen(
     bool noRegister) {
     opt = CodegenStackData::Allocate(this);
     opt->noRegister = noRegister;
-    opt->isRaster = true;
+    opt->isWorkGraph = true;
     auto disposeOpt = vstd::scope_exit([&] {
-        opt->isRaster = false;
+        opt->isWorkGraph = false;
         CodegenStackData::DeAllocate(std::move(opt));
     });
     vstd::StringBuilder codegenData;
@@ -470,12 +470,48 @@ CodegenResult CodegenUtility::WorkGraphCodegen(
     static_cast<void>(vstd::to_string(custom_mask));
     finalResult << '\n';
 
+    // Add work graph builtin template
+    finalResult << ReadInternalHLSLFile("work_graph"sv);
+    finalResult << "\n"sv;
+
     vstd::unordered_set<uint64_t> globalCallableMap;
-    for (const auto &node : work_graph.nodes()) {
-        CodegenWorkGraphNode(node, codegenData, globalCallableMap, false /* TODO: handle cbuffer-based arguments */);
+    const auto& nodes = work_graph.nodes();
+    const auto& entry_points = work_graph.entry_points();
+    vstd::unordered_set<uint32_t> entry_point_set;
+    for (auto ep : entry_points) {
+        entry_point_set.emplace(ep);
+    }
+    for (size_t i = 0; i < nodes.size(); ++i) {
+        bool is_entry_point = entry_point_set.contains(static_cast<uint32_t>(i));
+        CodegenWorkGraphNode(work_graph, i, is_entry_point, codegenData, globalCallableMap, false);
     }
 
-    LUISA_ASSERT(false, "unimplemented");
+    // Append the generated code for all nodes
+    finalResult << codegenData;
+
+    // Post-process properties (generates struct definitions)
+    PostprocessCodegenProperties(finalResult, false);
+
+    // Create the result
+    vstd::vector<Type const *> recordTypes;
+    recordTypes.reserve(nodes.size());
+    for (auto &&node : nodes) {
+        if (node.input_record_type != nullptr) {
+            recordTypes.push_back(node.input_record_type);
+        }
+    }
+    
+    auto result = CodegenResult(
+        std::move(finalResult),
+        std::move(opt->printer),
+        {},  // No properties for work graphs (resource bindings are handled differently)
+        opt->useTex2DBindless,
+        opt->useTex3DBindless,
+        opt->useBufferBindless,
+        immutableHeaderSize,
+        GetTypeMD5(recordTypes));
+
+    return result;
 }
 
 void CodegenUtility::CodegenFunction(Function func, vstd::StringBuilder &result, bool cbufferNonEmpty, bool codegen_self) {

src/backends/common/hlsl/codegen_utils/function_codegen.cpp

@@ -10,7 +10,6 @@
 #include <luisa/core/dynamic_module.h>
 #include <luisa/core/logging.h>
 #include <luisa/ast/external_function.h>
-#include "../constant_printer.h"
 
 // External declaration for shared variable from hlsl_codegen_util.cpp
 extern bool shown_buffer_warning;
@@ -1450,15 +1449,16 @@ void CodegenUtility::GetFunctionName(CallExpr const *expr, vstd::StringBuilder &
             return;
         case CallOp::WORK_GRAPH_OUTPUT: {
             LUISA_ASSERT(opt->isWorkGraph, "Work Graph Output can only be used in work graph nodes");
-            str << "_work_graph_output(";
-            str << "_work_graph_output_";
+            str << "_work_graph_output("sv;
+            str << "_work_graph_output_"sv;
             auto literal = static_cast<const LiteralExpr *>(args[0]);
             auto output_index = luisa::get<uint>(literal->value());
-            str << output_index << ", ";
+            vstd::to_string(output_index, str);
+            str << ", "sv;
             args[2]->accept(vis);
-            str << ", ";
+            str << ", "sv;
             args[3]->accept(vis);
-            str << ")";
+            str << ")"sv;
         } return;
         case CallOp::DDX: {
             if (opt->isRaster) {
@@ -1992,57 +1992,4 @@ o0=pixel(p,primId)"sv;
     // value assignment
 }
 
-void CodegenUtility::CodegenWorkGraphNode(const compute::detail::WorkGraphNode& node, vstd::StringBuilder &result, vstd::unordered_set<uint64_t>& callableMap, bool cbufferNonEmpty) {
-    auto codegenOneFunc = [&](Function func) {
-        auto constants = func.constants();
-        for (auto &&i : constants) {
-            vstd::StringBuilder constValueName;
-            if (!GetConstName(i.hash(), i, constValueName)) continue;
-            result << "static const "sv;
-            GetTypeName(*i.type(), result, Usage::READ);
-            result << ' ' << constValueName << " = "sv;
-            CodegenConstantPrinter printer{*this, result};
-            i.decode(printer);
-            result << ";\n"sv;
-        }
-#ifdef LUISA_ENABLE_IR
-        vstd::unordered_set<Variable> grad_vars;
-        // glob_variables_with_grad(func, grad_vars);
-#endif
-
-        opt->funcType = CodegenStackData::FuncType::Callable;
-        GetFunctionDecl(func, result);
-        result << "{\n"sv;
-        {
-
-            StringStateVisitor vis(func, result, this);
-            vis.sharedVariables = &opt->sharedVariable;
-            vis.VisitFunction(
-#ifdef LUISA_ENABLE_IR
-                grad_vars,
-#endif
-                func);
-        }
-        result << "}\n"sv;
-    };
-
-    auto callable = [&](auto &&callable, Function func) -> void {
-        for (auto &&i : func.custom_callables()) {
-            if (callableMap.emplace(i->hash()).second) {
-                callable(callable, i->function());
-            }
-        }
-        codegenOneFunc(func);
-    };
-
-    auto node_func = node.fn_builder->function();
-    for (auto &&i : node_func.custom_callables()) {
-        if (callableMap.emplace(i->hash()).second) {
-            callable(callable, i->function());
-        }
-    }
-
-    result << luisa::format("// node name: {}\n", node.name);
-    result << luisa::format("void {} (\n", node_func.name());
-}
 }// namespace lc::hlsl