wip(workgraph): add node type + appropriate attributes

Author: buggy213

include/luisa/dsl/work_graph/work_graph.h

@@ -24,8 +24,12 @@ struct WorkGraphNode {
     luisa::shared_ptr<const detail::FunctionBuilder> fn_builder;
 
     luisa::string name;
+    WorkGraphLaunchType node_type;
     const Type* input_record_type;
     luisa::vector<WorkGraphEdge> out_edges;
+    uint3 threadgroup_dim = uint3(1, 1, 1);
+    uint3 dispatch_dim = uint3(1, 1, 1);
+    luisa::optional<uint> dispatch_grid_member = luisa::nullopt;
     bool defined = false;
 };
 
@@ -60,7 +64,7 @@ class WorkGraphNodeOutput {
     uint _edge_index;
 };
 
-template<typename T>
+template<WorkGraphLaunchType NodeType, typename T>
 class WorkGraphNode {
 public:
     explicit WorkGraphNode(WorkGraphBuilder* builder, uint node_index) noexcept : _builder(builder), _node_index(node_index) {}
@@ -84,10 +88,8 @@ class WorkGraphNode {
         return WorkGraphNodeOutput<EdgeRecord>(_builder, _node_index, e.source_output_index);
     }
 
-    template<typename InputRecord, typename Def>
-    void define(const WorkGraphNodeKernel<InputRecord, Def>& kernel) noexcept {
-        static_assert(std::is_same_v<T, InputRecord>, "type mismatch between work graph node and its definition");
-
+    template<typename Def>
+    void define(const WorkGraphNodeKernel<T, Def>& kernel) noexcept {
         LUISA_ASSERT(!inner()->defined, "redefining node kernel is not allowed");
 
         // yoink the function builder, make sure type of input record matches what we were declared with
@@ -111,6 +113,35 @@ class WorkGraphNode {
         edge->dest = inner()->index;
     }
 
+    // specifying NumThreads (threadgroup size) is not allowed for per-thread node launch
+    void set_threadgroup_size(uint3 size) const requires (NodeType != WorkGraphLaunchType::THREAD) {
+        inner()->threadgroup_dim = size;
+    }
+
+    // for broadcasting nodes, this is either the static size of a dispatch, or the max size
+    // (if its dynamically sized using SV_DispatchGrid)
+    void set_dispatch_size(uint3 size) const requires (NodeType == WorkGraphLaunchType::BROADCASTING) {
+        inner()->dispatch_dim = size;
+    }
+
+    // sets field in InputRecord type to be marked with SV_DispatchGrid
+    void set_dispatch_grid_field(uint index) const requires (NodeType == WorkGraphLaunchType::BROADCASTING) {
+        auto t = Type::of<T>();
+        auto members = t->members();
+
+        LUISA_ASSERT(index < members.size(), "dispatch grid field index out of bounds");
+
+        auto member_type = members[index];
+        LUISA_ASSERT(
+            member_type == Type::of<uint>() ||
+            member_type == Type::of<uint2>() ||
+            member_type == Type::of<uint3>(),
+            "dispatch grid field must be uint, uint2, or uint3"
+        );
+
+        inner()->dispatch_grid_member = index;
+    }
+
 private:
     // invalidated by modifying graph
     [[nodiscard]] detail::WorkGraphNode* inner() const noexcept { return &detail::index_to_node(_builder, _node_index); }
@@ -139,8 +170,8 @@ class WorkGraph {
 class WorkGraphBuilder {
 public:
 
-    template<typename InputRecord>
-    WorkGraphNode<InputRecord> add_node(luisa::string name) noexcept {
+    template<WorkGraphLaunchType NodeType, typename InputRecord>
+    WorkGraphNode<NodeType, InputRecord> add_node(luisa::string name) noexcept {
         const Type* input_record_type;
         if constexpr (std::is_same_v<InputRecord, WorkGraphEmptyRecord>) {
             input_record_type = nullptr;
@@ -155,10 +186,11 @@ class WorkGraphBuilder {
             nullptr,
 
             std::move(name),
+            NodeType,
             input_record_type
         );
 
-        return WorkGraphNode<InputRecord>(this, node_index);
+        return WorkGraphNode<NodeType, InputRecord>(this, node_index);
     }
 
     LUISA_DSL_API WorkGraph build() noexcept;

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

@@ -59,22 +59,47 @@ void CodegenUtility::GenerateNodeInputDecl(
 
 // Helper to generate node shader attributes
 void CodegenUtility::GenerateNodeShaderAttributes(
-    bool is_entry_point,
-    luisa::string_view node_name,
+    const luisa::compute::detail::WorkGraphNode &node,
     vstd::StringBuilder &result) {
 
     // [Shader("node")] attribute (required for all work graph nodes)
     result << "[Shader(\"node\")]\n"sv;
 
-    // Entry point nodes need [NodeLaunch("broadcasting")]
-    if (is_entry_point) {
-        result << "[NodeLaunch(\"broadcasting\")]\n"sv;
+    switch (node.node_type) {
+        case WorkGraphLaunchType::BROADCASTING: {
+            result << "[NodeLaunch(\"broadcasting\")]\n"sv;
+
+            luisa::string threadgroup_dim = luisa::format(
+                "[NumThreads({}, {}, {})]",
+                node.threadgroup_dim.x, node.threadgroup_dim.y, node.threadgroup_dim.z
+            );
+            result << threadgroup_dim << '\n';
+
+            luisa::string dispatch_properties;
+            if (node.dispatch_grid_member) {
+                dispatch_properties = luisa::format(
+                    "[NodeMaxDispatchGrid({}, {}, {})]",
+                    node.dispatch_dim.x, node.dispatch_dim.y, node.dispatch_dim.z
+                );
+            }
+            else {
+                dispatch_properties = luisa::format(
+                    "[NodeDispatchGrid({}, {}, {})]",
+                    node.dispatch_dim.x, node.dispatch_dim.y, node.dispatch_dim.z
+                );
+            }
+            result << dispatch_properties << '\n';
+        } break;
+        case WorkGraphLaunchType::THREAD: {
+            result << "[NodeLaunch(\"thread\")]\n"sv;
+        } break;
     }
 
     // Add comment with node name for debugging
-    if (!node_name.empty()) {
-        result << "// node name: "sv << node_name << "\n"sv;
+    if (!node.name.empty()) {
+        result << "// node name: "sv << node.name << "\n"sv;
     }
+
 }
 
 // Helper to generate work graph node function signature
@@ -270,7 +295,7 @@ void CodegenUtility::CodegenWorkGraphNode(const WorkGraph &work_graph, size_t no
     }
 
     // Generate node shader attributes using helper
-    GenerateNodeShaderAttributes(is_entry_point, node.name, result);
+    GenerateNodeShaderAttributes(node, result);
 
     // Generate node function signature
     // use actual name from frontend here, rather than custom_<i>, since node names are meaningful

src/backends/common/hlsl/hlsl_codegen.h

@@ -73,7 +73,7 @@ class CodegenUtility {
     static void GenerateMaxRecordsAttribute(uint max_records, vstd::StringBuilder &result);
     void GenerateNodeOutputDecl(const Type *record_type, uint max_records, luisa::string_view var_name_prefix, int output_index, vstd::StringBuilder &result);
     void GenerateNodeInputDecl(const Type *record_type, luisa::string_view var_name, vstd::StringBuilder &result);
-    static void GenerateNodeShaderAttributes(bool is_entry_point, luisa::string_view node_name, vstd::StringBuilder &result);
+    static void GenerateNodeShaderAttributes(const luisa::compute::detail::WorkGraphNode& node, vstd::StringBuilder &result);
     void GenerateNodeFunctionSignature(Function node_func, const luisa::compute::detail::WorkGraphNode &node, const luisa::vector<luisa::compute::detail::WorkGraphNode> &all_nodes, vstd::StringBuilder &result);
     static void GenerateWorkGraphOutputCall(int output_index, luisa::string_view record_var_name, vstd::StringBuilder &result);
     void GenerateRecordStructDef(const Type *record_type, vstd::StringBuilder &result);

src/backends/common/hlsl/test/test_work_graph_codegen.cpp

@@ -42,7 +42,7 @@ WorkGraph create_simple_entry_to_consumer() {
     WorkGraphBuilder wg;
 
     // Entry point node (always uses WorkGraphEmptyRecord)
-    auto entry = wg.add_node<WorkGraphEmptyRecord>("entry_node");
+    auto entry = wg.add_node<WorkGraphLaunchType::BROADCASTING, WorkGraphEmptyRecord>("entry_node");
     auto entry_output = entry.output<SimpleRecord>(16);
 
     WorkGraphNodeKernel entry_kernel = [&]() {
@@ -53,7 +53,7 @@ WorkGraph create_simple_entry_to_consumer() {
     entry.define(entry_kernel);
 
     // Consumer node (receives SimpleRecord)
-    auto consumer = wg.add_node<SimpleRecord>("consumer_node");
+    auto consumer = wg.add_node<WorkGraphLaunchType::THREAD, SimpleRecord>("consumer_node");
     WorkGraphNodeKernel consumer_kernel = [&](Var<SimpleRecord> input) {
         auto val = input->value;
         (void)val;
@@ -71,7 +71,7 @@ WorkGraph create_multi_output_node() {
     WorkGraphBuilder wg;
 
     // Entry with multiple outputs
-    auto entry = wg.add_node<WorkGraphEmptyRecord>("multi_out_entry");
+    auto entry = wg.add_node<WorkGraphLaunchType::BROADCASTING, WorkGraphEmptyRecord>("multi_out_entry");
     auto output_a = entry.output<SimpleRecord>(8);
     auto output_b = entry.output<ComplexRecord>(4);
 
@@ -90,15 +90,15 @@ WorkGraph create_multi_output_node() {
     entry.define(entry_kernel);
 
     // Consumer A (receives SimpleRecord)
-    auto consumer_a = wg.add_node<SimpleRecord>("consumer_a");
+    auto consumer_a = wg.add_node<WorkGraphLaunchType::THREAD, SimpleRecord>("consumer_a");
     WorkGraphNodeKernel consumer_a_kernel = [&](Var<SimpleRecord> input) {
         auto val = input->value;
         (void)val;
     };
     consumer_a.define(consumer_a_kernel);
 
     // Consumer B (receives ComplexRecord)
-    auto consumer_b = wg.add_node<ComplexRecord>("consumer_b");
+    auto consumer_b = wg.add_node<WorkGraphLaunchType::THREAD, ComplexRecord>("consumer_b");
     WorkGraphNodeKernel consumer_b_kernel = [&](Var<ComplexRecord> input) {
         auto id = input->id;
         auto pos = input->position;
@@ -119,7 +119,7 @@ WorkGraph create_chained_nodes() {
     WorkGraphBuilder wg;
 
     // Node A: Entry point
-    auto node_a = wg.add_node<WorkGraphEmptyRecord>("node_a");
+    auto node_a = wg.add_node<WorkGraphLaunchType::BROADCASTING, WorkGraphEmptyRecord>("node_a");
     auto output_a = node_a.output<SimpleRecord>(32);
 
     WorkGraphNodeKernel node_a_kernel = [&]() {
@@ -130,7 +130,7 @@ WorkGraph create_chained_nodes() {
     node_a.define(node_a_kernel);
 
     // Node B: Middle node (processes and forwards)
-    auto node_b = wg.add_node<SimpleRecord>("node_b");
+    auto node_b = wg.add_node<WorkGraphLaunchType::THREAD, SimpleRecord>("node_b");
     auto output_b = node_b.output<SimpleRecord>(16);
 
     WorkGraphNodeKernel node_b_kernel = [&](Var<SimpleRecord> input) {
@@ -141,7 +141,7 @@ WorkGraph create_chained_nodes() {
     node_b.define(node_b_kernel);
 
     // Node C: Final node
-    auto node_c = wg.add_node<SimpleRecord>("node_c");
+    auto node_c = wg.add_node<WorkGraphLaunchType::THREAD, SimpleRecord>("node_c");
     WorkGraphNodeKernel node_c_kernel = [&](Var<SimpleRecord> input) {
         auto val = input->value;
         (void)val;
@@ -160,7 +160,7 @@ WorkGraph create_terminal_entry_node() {
     WorkGraphBuilder wg;
 
     // Entry point that doesn't output anything
-    auto entry = wg.add_node<WorkGraphEmptyRecord>("terminal_entry");
+    auto entry = wg.add_node<WorkGraphLaunchType::BROADCASTING, WorkGraphEmptyRecord>("terminal_entry");
 
     WorkGraphNodeKernel entry_kernel = [&]() {
         // Just do some work, no outputs

src/tests/test_work_graph.cpp

@@ -19,7 +19,8 @@ LUISA_STRUCT(ConsumerRecord, datum) {};
 WorkGraph describe_work_graph() {
     WorkGraphBuilder work_graph;
 
-    auto producer = work_graph.add_node<WorkGraphEmptyRecord>("producer");
+    auto producer = work_graph.add_node<WorkGraphLaunchType::BROADCASTING, WorkGraphEmptyRecord>("producer");
+    producer.set_threadgroup_size({64, 1, 1});
     auto producer_output = producer.output<ConsumerRecord>(16);
     WorkGraphNodeKernel producer_kernel = [&]() {
         Var<ConsumerRecord> out;
@@ -28,7 +29,7 @@ WorkGraph describe_work_graph() {
     producer.define(producer_kernel);
 
 
-    auto consumer = work_graph.add_node<ConsumerRecord>("consumer");
+    auto consumer = work_graph.add_node<WorkGraphLaunchType::THREAD, ConsumerRecord>("consumer");
     WorkGraphNodeKernel consumer_kernel = [&](Var<ConsumerRecord> input) {
         // do work
     };