1	# A work partition is independent of its element representation and ownership.
2	# Storage remains in the leaf payload; offset/span describe its logical range.
3	import Base
4	import ./storage_buffer.bend as Storage
5	
6	type Tree<ownership,-Block: Kind(ownership),-Details: Data> is Kind(ownership):
7	  Leaf{data: Block,offset: U32,span: U32,details: Details}
8	  Fork{left: Tree<ownership,Block,Details>,right: Tree<ownership,Block,Details>}
9	
10	# Each binary fork gives both children independent owners; its leaf callback
11	# finishes all local stages before returning. The input split sees the left
12	# child's range, so it can copy only what that child reads. Representation
13	# adapters supply direct templates, so Array and certified Buffer share it.
14	# The last split calls its two leaves directly: a leaf that neither forks nor
15	# calls a closure is then its own task, which the runtime keeps for the work
16	# turn; a leaf reached through the forking recursion would run in the grow
17	# turn on one worker per row.
18	def run(~Input: Type,~Weights: Type,~Bias: Type,~Output: Type,~Context: Data,
19	  ~split_input: Context -> U32 -> U32 -> Input -> Input & Input,~clone_weights: Weights -> Weights & Weights,~clone_bias: Bias -> Bias & Bias,
20	  ~leaf: Context -> U32 -> U32 -> Input -> Weights -> Bias -> Tree<&1,Output,U32>,
21	  depth: Nat,+context: Context,+offset: U32,+span: U32,input: Input,weights: Weights,bias: Bias) -> Tree<&1,Output,U32>:
22	  match depth:
23	    case 0n: leaf(context,offset,span,input,weights,bias)
24	    case 1n:
25	      +half = (span / 16 * 8 : U32)
26	      Storage.bind_pair(~Input,~Input,~Tree<&1,Output,U32>,split_input(context,offset,half,input),left_input => right_input =>
27	        Storage.bind_pair(~Weights,~Weights,~Tree<&1,Output,U32>,clone_weights(weights),left_weights => right_weights =>
28	          Storage.bind_pair(~Bias,~Bias,~Tree<&1,Output,U32>,clone_bias(bias),left_bias => right_bias =>
29	            left right = leaf(context,offset,half,left_input,left_weights,left_bias) leaf(context,(offset + half : U32),(span - half : U32),right_input,right_weights,right_bias)
30	            Fork{left,right})))
31	    case 2n+ +rest:
32	      +half = (span / 16 * 8 : U32)
33	      Storage.bind_pair(~Input,~Input,~Tree<&1,Output,U32>,split_input(context,offset,half,input),left_input => right_input =>
34	        Storage.bind_pair(~Weights,~Weights,~Tree<&1,Output,U32>,clone_weights(weights),left_weights => right_weights =>
35	          Storage.bind_pair(~Bias,~Bias,~Tree<&1,Output,U32>,clone_bias(bias),left_bias => right_bias =>
36	            left right = run(~Input,~Weights,~Bias,~Output,~Context,~split_input,~clone_weights,~clone_bias,~leaf,1n+rest,context,offset,half,left_input,left_weights,left_bias) run(~Input,~Weights,~Bias,~Output,~Context,~split_input,~clone_weights,~clone_bias,~leaf,1n+rest,context,(offset + half : U32),(span - half : U32),right_input,right_weights,right_bias)
37	            Fork{left,right})))
38	
39	def gather(~Output: Type,~copy: U32 -> U32 -> U32 -> U32 -> Output -> Output -> Output,
40	  partitions: Tree<&1,Output,U32>,output: Output,+positions: U32) -> Output:
41	  match partitions:
42	    case Leaf{input,offset,span,rows}: copy(rows,offset,positions,span,input,output)
43	    case Fork{left,right}: gather(~Output,~copy,right,gather(~Output,~copy,left,output,positions),positions)