nx_task_graph.nx source

1// nx_task_graph.nx -- DAG of tasks with explicit dependency edges. 2// 3// Built on nx_thread_pool: each node is a (fn_ptr, ctx) task plus 4// dependency edges to/from other nodes. Scheduling is topo-sort 5// via atomic counters: when a task finishes, it atomic-decrements 6// each successor's `predecessors_done` counter; any successor whose 7// counter reaches `n_predecessors` is now ready and gets submitted 8// to the pool. Initial frontier (nodes with zero predecessors) is 9// submitted by `nx_graph_run`. 10// 11// Why this matters for SSS-class workloads: 12// * ML inference: layer N can't run before layer N-1 finishes. 13// Task graph naturally expresses this without manual barriers. 14// * Render pipelines: vertex -> raster -> pixel -> tonemap 15// stages with parallelism inside each stage. 16// * Build systems: target depends on N source files; ninja is 17// basically a task-graph engine. 18// * MapReduce: map tasks -> shuffle -> reduce tasks. 19// 20// MVP shape -- 8 successors per node ceiling (fixed-size inline), 21// no cycle detection (caller's responsibility), no priorities. 22// All three are next-evolution concerns. 23// 24// Composes against: [[nx_thread_pool_shared_queue]] (executor), 25// [[atomic_intrinsics_real_amo]] (predecessor counters), 26// [[fn_ptr_indirect_call]] (typed task dispatch), 27// [[vyukov_mpmc_channel]] (pool's task queue). 28 29// nx_safety_envelope: 30// intended_use: AUTO_APPLIED -- primitive-specific tuning queued 31// sil_target: SIL1 32// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail] 33// verdict: NOT_YET_EVALUATED 34 35import "nx_syscalls.nx" 36import "nx_atom.nx" 37import "nx_thread_pool.nx" 38const NX_MAGIC_2000000000: i64 = 2000000000 39 40const NX_GRAPH_MAX_SUCCS_PER_NODE: i64 = 8 41 42const NX_GRAPH_STATE_PENDING: i64 = 0 43const NX_GRAPH_STATE_READY: i64 = 1 44const NX_GRAPH_STATE_RUNNING: i64 = 2 45const NX_GRAPH_STATE_DONE: i64 = 3 46 47// Layout: 16 i64 fields = 128 bytes per node. succ_0..succ_7 are 48// contiguous so we can read them via pointer arithmetic from a base. 49struct NxGraphNode { 50 graph_ptr: i64, // back-pointer to NxTaskGraph 51 task_fn: func(i64) -> i64, // user task body 52 user_ctx: i64, // arg passed to task_fn 53 n_predecessors: i64, // set during build 54 predecessors_done: i64, // atomic; bumped when a pred completes 55 n_successors: i64, 56 state: i64, 57 _pad0: i64, 58 succ_0: i64, 59 succ_1: i64, 60 succ_2: i64, 61 succ_3: i64, 62 succ_4: i64, 63 succ_5: i64, 64 succ_6: i64, 65 succ_7: i64, 66} 67 68const NX_GRAPH_NODE_BYTES: i64 = 128 69const NX_GRAPH_OFF_PRED_DONE: i64 = 32 // offset of predecessors_done within NxGraphNode 70const NX_GRAPH_OFF_SUCC_0: i64 = 64 // offset of succ_0 within NxGraphNode 71 72struct NxTaskGraph { 73 pool_ptr: i64, 74 nodes_base: i64, // *NxGraphNode 75 n_nodes: i64, 76 max_nodes: i64, 77 completed_count: i64, // atomic; bumped per task finish 78 _pad: i64, 79} 80 81const NX_GRAPH_OFF_COMPLETED: i64 = 32 82 83// Compute address of node[idx]. 84func _nx_graph_node_at(g: *NxTaskGraph, idx: i64) -> *NxGraphNode { 85 return (g.nodes_base + idx * NX_GRAPH_NODE_BYTES) as *NxGraphNode 86} 87 88// Write succ_K via pointer arithmetic. 89func _nx_graph_set_succ(node: *NxGraphNode, k: i64, succ_idx: i64) -> i64 { 90 let base: *i64 = ((node as i64) + NX_GRAPH_OFF_SUCC_0) as *i64 91 base[k] = succ_idx 92 return 0 93} 94 95// Read succ_K via pointer arithmetic. 96func _nx_graph_get_succ(node: *NxGraphNode, k: i64) -> i64 { 97 let base: *i64 = ((node as i64) + NX_GRAPH_OFF_SUCC_0) as *i64 98 return base[k] 99} 100 101// Wrapper that every node runs as its pool task. ctx is the node's 102// address. Runs user fn, then signals successors atomically, then 103// submits any newly-ready successor. Finally bumps completed_count. 104func _nx_graph_run_node(ctx: i64) -> i64 { 105 let node: *NxGraphNode = ctx as *NxGraphNode 106 let g: *NxTaskGraph = node.graph_ptr as *NxTaskGraph 107 108 node.state = NX_GRAPH_STATE_RUNNING 109 let fp: func(i64) -> i64 = node.task_fn 110 fp(node.user_ctx) 111 node.state = NX_GRAPH_STATE_DONE 112 113 // Signal successors. 114 let pool: *NxThreadPool = g.pool_ptr as *NxThreadPool 115 var k: i64 = 0 116 while k < node.n_successors { 117 let s_idx: i64 = _nx_graph_get_succ(node, k) 118 let succ: *NxGraphNode = _nx_graph_node_at(g, s_idx) 119 let pd_addr: *i64 = ((succ as i64) + NX_GRAPH_OFF_PRED_DONE) as *i64 120 let prior: i64 = nx_atom_faa_i64(pd_addr, 1, NX_MO_SEQ_CST) 121 let now_done: i64 = prior + 1 122 if now_done == succ.n_predecessors { 123 // Successor's deps all satisfied -- submit it. 124 nx_pool_submit(pool, _nx_graph_run_node, succ as i64) 125 } 126 k = k + 1 127 } 128 129 // Bump graph completed counter. 130 let cc_addr: *i64 = ((g as i64) + NX_GRAPH_OFF_COMPLETED) as *i64 131 nx_atom_faa_i64(cc_addr, 1, NX_MO_SEQ_CST) 132 return 0 133} 134 135// ---- public API -------------------------------------------------- 136 137// Create a new task graph backed by `pool`, sized for up to 138// max_nodes nodes. Returns NULL-ish (-1 cast) on OOM. 139func nx_graph_new(pool: *NxThreadPool, max_nodes: i64) -> *NxTaskGraph { 140 let raw: *u8 = sys_mmap(128) 141 let g: *NxTaskGraph = raw as *NxTaskGraph 142 let nodes_raw: *u8 = sys_mmap(max_nodes * NX_GRAPH_NODE_BYTES) 143 g.pool_ptr = pool as i64 144 g.nodes_base = nodes_raw as i64 145 g.n_nodes = 0 146 g.max_nodes = max_nodes 147 g.completed_count = 0 148 return g 149} 150 151// Append a new node to the graph. Returns its index (>=0) or -1 on 152// capacity exhaustion. Node starts with zero predecessors and zero 153// successors -- caller uses nx_graph_add_edge to wire dependencies. 154func nx_graph_add_node(g: *NxTaskGraph, fn: func(i64) -> i64, ctx: i64) -> i64 { 155 if g.n_nodes >= g.max_nodes { return -1 } 156 let idx: i64 = g.n_nodes 157 let node: *NxGraphNode = _nx_graph_node_at(g, idx) 158 node.graph_ptr = g as i64 159 node.task_fn = fn 160 node.user_ctx = ctx 161 node.n_predecessors = 0 162 node.predecessors_done = 0 163 node.n_successors = 0 164 node.state = NX_GRAPH_STATE_PENDING 165 g.n_nodes = idx + 1 166 return idx 167} 168 169// Add a from -> to dependency edge. `to` cannot run until `from` 170// completes. Returns 0 on success, -1 if `from` has already hit 171// the per-node successor ceiling. 172func nx_graph_add_edge(g: *NxTaskGraph, from_idx: i64, to_idx: i64) -> i64 { 173 let from_node: *NxGraphNode = _nx_graph_node_at(g, from_idx) 174 if from_node.n_successors >= NX_GRAPH_MAX_SUCCS_PER_NODE { return -1 } 175 _nx_graph_set_succ(from_node, from_node.n_successors, to_idx) 176 from_node.n_successors = from_node.n_successors + 1 177 let to_node: *NxGraphNode = _nx_graph_node_at(g, to_idx) 178 to_node.n_predecessors = to_node.n_predecessors + 1 179 return 0 180} 181 182// Execute the graph: submit all zero-predecessor nodes to the pool, 183// then spin-wait until completed_count equals n_nodes. Returns 0 184// on success, -1 on timeout. 185func nx_graph_run(g: *NxTaskGraph) -> i64 { 186 // Submit initial frontier. 187 var i: i64 = 0 188 while i < g.n_nodes { 189 let node: *NxGraphNode = _nx_graph_node_at(g, i) 190 if node.n_predecessors == 0 { 191 let pool: *NxThreadPool = g.pool_ptr as *NxThreadPool 192 nx_pool_submit(pool, _nx_graph_run_node, node as i64) 193 } 194 i = i + 1 195 } 196 197 // Spin until all nodes have signalled completion. Yield between 198 // probes -- single-core qemu otherwise starves worker threads. 199 let cc_addr: *i64 = ((g as i64) + NX_GRAPH_OFF_COMPLETED) as *i64 200 var spins: i64 = 0 201 while nx_atom_load_i64(cc_addr, NX_MO_SEQ_CST) < g.n_nodes { 202 nx_thread_yield() 203 spins = spins + 1 204 if spins > NX_MAGIC_2000000000 { return -1 } 205 } 206 return 0 207} 208 209func nx_graph_n_completed(g: *NxTaskGraph) -> i64 { 210 let cc_addr: *i64 = ((g as i64) + NX_GRAPH_OFF_COMPLETED) as *i64 211 return nx_atom_load_i64(cc_addr, NX_MO_SEQ_CST) 212} 213 214func nx_graph_node_state(g: *NxTaskGraph, idx: i64) -> i64 { 215 let node: *NxGraphNode = _nx_graph_node_at(g, idx) 216 return node.state 217} 218 219// ---- self-test --------------------------------------------------- 220 221// Each test task FAA-bumps a counter at the address passed in ctx. 222func _graph_self_test_task(ctx: i64) -> i64 { 223 let p: *i64 = ctx as *i64 224 nx_atom_faa_i64(p, 1, NX_MO_SEQ_CST) 225 return 0 226} 227 228func main() -> i64 { 229 let pool: *NxThreadPool = nx_pool_new(2, 32) 230 231 // Build a 5-node graph and verify all run. 232 let g: *NxTaskGraph = nx_graph_new(pool, 16) 233 let counter_raw: *u8 = sys_mmap(16) 234 let counter: *i64 = counter_raw as *i64 235 *counter = 0 236 let counter_addr: i64 = counter as i64 237 238 var i: i64 = 0 239 while i < 5 { 240 nx_graph_add_node(g, _graph_self_test_task, counter_addr) 241 i = i + 1 242 } 243 // Chain: 0 -> 1 -> 2 -> 3 -> 4 (pure linear) 244 nx_graph_add_edge(g, 0, 1) 245 nx_graph_add_edge(g, 1, 2) 246 nx_graph_add_edge(g, 2, 3) 247 nx_graph_add_edge(g, 3, 4) 248 249 if nx_graph_run(g) != 0 { return __syscall(93, 1, 0, 0, 0, 0, 0) } 250 if nx_graph_n_completed(g) != 5 { return __syscall(93, 2, 0, 0, 0, 0, 0) } 251 if *counter != 5 { return __syscall(93, 3, 0, 0, 0, 0, 0) } 252 253 nx_pool_shutdown(pool) 254 return 0 255}