1	# Structural cost of the 4x8 pipeline. Only guarded dimensions enter here.
2	# Packing and GEMM run in each leaf; reordering and final gather are serial.
3	import Base
4	import ./kernel_cost.bend as Cost
5	import ./kernel_window.bend as KernelWindow
6	import ./math_activation.bend as Activations
7	import ./kernel_pipeline.bend as Pipeline
8	
9	type Owners is Data:
10	  Owners{input_capacity: U32,bias_capacity: U32}
11	
12	# Image and stencil extents that decide which input rows a fork copies.
13	type Geometry is Data:
14	  Geometry{channels: U32,height: U32,width: U32,kernel: U32,stride: U32,padding: U32,output_width: U32}
15	
16	type Work is Data:
17	  Work{positions: U32,inner: U32,rows: U32,owners: Owners,geometry: Geometry,activation: Activations.Activation}
18	
19	# The input rows a partition node holds and the capacity of their owner.
20	type InputRows is Data:
21	  InputRows{first: U32,rows: U32,capacity: Nat}
22	
23	# Allocation and moved/initialized payload are distinct for a copied band.
24	type ForkInput is Data:
25	  ForkInput{bytes: Nat,transfer: Nat,left: InputRows}
26	
27	# critical_ns is the span (the costliest root-to-leaf path) and work_ns the sum
28	# over every leaf and fork; together they bound time on a given worker count.
29	type Movement is Data:
30	  Movement{packed: Nat,leaves: Nat,copies: Nat,critical_ns: F32,work_ns: F32,minimum_leaf_ns: F32}
31	
32	type Plan is Data:
33	  Plan{packed: Bool,split_depth: U32,workspace_bytes: Nat,estimated_ns: F32}
34	
35	def capacity_loop(remaining: Nat,+value: U32,+capacity: Nat,zero: Bool) -> Nat:
36	  match remaining zero:
37	    case 0n zero: capacity
38	    case 1n+rest True{}: capacity
39	    case 1n+rest False{}:
40	      +next = U32.shr(value)
41	      capacity_loop(rest,next,Nat.double(capacity),U32.is_zero(next))
42	
43	def capacity(+elements: U32) -> Nat:
44	  +value = (U32.max(1,elements) - 1 : U32)
45	  capacity_loop(32n,value,1n,U32.is_zero(value))
46	
47	# A byte count as an F32: its high word times 2^32 plus its low word, which
48	# F32.from_nat reads as the count modulo 2^32.
49	def high_word(halvings: Nat,+bytes: Nat) -> Nat:
50	  match halvings:
51	    case 0n: bytes
52	    case 1n+rest: high_word(rest,Nat.div(bytes,2n))
53	
54	def payload_float(+bytes: Nat) -> F32:
55	  (F32.from_nat(high_word(32n,bytes)) * 4294967296.0 + F32.from_nat(bytes) : F32)
56	
57	def weight_bytes(work: Work) -> Nat:
58	  Work{positions,inner,rows,owners,geometry,activation} = work
59	  Nat.mul(32n,capacity(((rows + 7) / 8 * inner : U32)))
60	
61	# The weight reorder writes complete eight-channel tiles, including zero lanes.
62	def reorder_ns(+work: Work,config: Cost.Configuration) -> F32:
63	  Work{positions,+inner,+rows,owners,geometry,activation} = work
64	  Cost.Configuration{mac,byte_ns,operation_ns,grain,wake} = config
65	  ((U32.to_f32(((rows + 7) / 8 * 8 : U32)) * U32.to_f32(inner) * 4.0 + payload_float(weight_bytes(work))) * byte_ns : F32)
66	
67	# Each fork clones the reordered weights and bias, and splits the input rows.
68	def fork_bytes(+work: Work,input: Nat) -> Nat:
69	  Work{positions,inner,rows,Owners{input_capacity,bias},geometry,activation} = work
70	  Nat.add(weight_bytes(work),Nat.add(input,Nat.mul(4n,U32.to_nat(U32.max(2,bias)))))
71	
72	def clone_bytes(+window: InputRows) -> Nat:
73	  InputRows{first,rows,+capacity} = window
74	  Nat.mul(4n,Nat.max(2n,capacity))
75	
76	def band_rows(+band: KernelWindow.Band,+channels: U32,+width: U32) -> InputRows:
77	  KernelWindow.Band{+first,+rows} = band
78	  InputRows{first,rows,capacity((channels * rows * width : U32))}
79	
80	def band_bytes(+band: KernelWindow.Band,+channels: U32,+width: U32) -> Nat:
81	  KernelWindow.Band{first,+rows} = band
82	  Nat.mul(4n,Nat.max(2n,capacity((channels * rows * width : U32))))
83	
84	# copied initializes the new owner and copies every logical band value.
85	# Array.clone instead copies the complete allocation, without initialization.
86	def input_transfer(copy: Bool,band: KernelWindow.Band,window: InputRows,channels: U32,width: U32) -> Nat:
87	  match copy band:
88	    case False{} band: clone_bytes(window)
89	    case True{} KernelWindow.Band{first,rows}:
90	      +count = (channels * rows * width : U32)
91	      Nat.mul(4n,Nat.add(U32.to_nat(count),capacity(count)))
92	
93	# As in the pipelines: a checked band over aligned rows is copied and the
94	# second half keeps the parent; otherwise the parent owner is cloned.
95	def input_bytes(copy: Bool,+band: KernelWindow.Band,+window: InputRows,+channels: U32,+width: U32) -> Nat:
96	  match copy:
97	    case False{}: clone_bytes(window)
98	    case True{}: band_bytes(band,channels,width)
99	
100	def left_rows(copy: Bool,+band: KernelWindow.Band,+window: InputRows,+channels: U32,+width: U32) -> InputRows:
101	  match copy:
102	    case False{}: window
103	    case True{}: band_rows(band,channels,width)
104	
105	def fork_window(+offset: U32,+half: U32,+window: InputRows,+geometry: Geometry) -> ForkInput:
106	  InputRows{+first,+count,+capacity} = window
107	  Geometry{+channels,+height,+width,+kernel,+stride,+padding,+output_width} = geometry
108	  +band = KernelWindow.band(first,count,height,kernel,stride,padding,output_width,offset,half)
109	  +copy = KernelWindow.copies_band(first,count,height,kernel,stride,padding,output_width,offset,half,band)
110	  +parent = {InputRows{first,count,capacity} : InputRows}
111	  ForkInput{input_bytes(copy,band,parent,channels,width),input_transfer(copy,band,parent,channels,width),left_rows(copy,band,parent,channels,width)}
112	
113	def forked_bytes(+fork: ForkInput) -> Nat:
114	  ForkInput{+bytes,transfer,left} = fork
115	  bytes
116	
117	def forked_transfer(+fork: ForkInput) -> Nat:
118	  ForkInput{bytes,+transfer,left} = fork
119	  transfer
120	
121	def forked_left(+fork: ForkInput) -> InputRows:
122	  ForkInput{bytes,transfer,+left} = fork
123	  left
124	
125	def root_rows(+work: Work) -> InputRows:
126	  Work{positions,inner,rows,Owners{+input,bias},Geometry{channels,+height,width,kernel,stride,padding,output_width},activation} = work
127	  InputRows{0,height,U32.to_nat(input)}
128	
129	# The activation pass reads and writes each logical output in place; the
130	# identity skips the pass.
131	def activation_ns(+activation: Activations.Activation,+config: Cost.Configuration) -> F32:
132	  Bool.pick(F32,Activations.is_identity(activation),0.0,Cost.element_ns(4.0,Activations.operations(activation),config))
133	
134	def leaf_cost(span: U32,+count: U32,work: Work,+config: Cost.Configuration) -> Movement:
135	  Work{positions,+inner,+rows,owners,geometry,activation} = work
136	  Cost.Configuration{mac_ns,+byte_ns,operation_ns,grain,wake} = config
137	  +span = span
138	  +packed = Nat.mul(32n,capacity(U32.max(Pipeline.block_tiles(),inner)))
139	  +output_capacity = capacity((rows * count : U32))
140	  +output = Nat.mul(4n,Nat.max(2n,output_capacity))
141	  # Padded lanes still execute every MAC, including absent rows of the last unit.
142	  +core = (U32.to_f32(((rows + 3) / 4 * 4 : U32)) * U32.to_f32(span) * U32.to_f32(inner) * mac_ns : F32)
143	  # Initializing each allocation and transferring logical data are separate work.
144	  +packing = ((U32.to_f32(span) * U32.to_f32(inner) * 4.0 + payload_float(packed)) * byte_ns : F32)
145	  +stores = ((U32.to_f32(rows) * U32.to_f32(count) * 4.0 + payload_float(Nat.mul(4n,output_capacity))) * byte_ns : F32)
146	  +finish = (U32.to_f32(rows) * U32.to_f32(count) * activation_ns(activation,config) : F32)
147	  +time = (core + packing + stores + finish : F32)
148	  Movement{packed,output,0n,time,time,time}
149	
150	def merge(bytes: Nat,transfer: Nat,byte_ns: F32,left: Movement,right: Movement) -> Movement:
151	  Movement{lp,lo,lc,lt,lw,lmin} = left
152	  Movement{rp,ro,rc,rt,rw,rmin} = right
153	  +copy = (payload_float(transfer) * byte_ns : F32)
154	  Movement{Nat.add(lp,rp),Nat.add(lo,ro),Nat.add(bytes,Nat.add(lc,rc)),
155	    (copy + F32.max(lt,rt) : F32),(copy + lw + rw : F32),F32.min(lmin,rmin)}
156	
157	def movement_with(~Context: Data,~fork: @+offset: U32 -> @+half: U32 -> @+window: InputRows -> @+geometry: Context -> ForkInput,
158	  depth: Nat,+offset: U32,+span: U32,+window: InputRows,+geometry: Context,+work: Work,+config: Cost.Configuration) -> Movement:
159	  match depth:
160	    case 0n:
161	      Work{+positions,inner,rows,owners,geometry,activation} = work
162	      leaf_cost(span,U32.min(span,(positions - offset : U32)),work,config)
163	    case 1n+ +rest:
164	      Cost.Configuration{mac,byte_ns,operation_ns,grain,wake} = config
165	      +half = (span / 16 * 8 : U32)
166	      +input = fork(offset,half,window,geometry)
167	      merge(fork_bytes(work,forked_bytes(input)),fork_bytes(work,forked_transfer(input)),byte_ns,
168	        movement_with(~Context,~fork,rest,offset,half,forked_left(input),geometry,work,config),
169	        movement_with(~Context,~fork,rest,(offset + half : U32),(span - half : U32),window,geometry,work,config))
170	
171	def movement(depth: Nat,offset: U32,span: U32,window: InputRows,+work: Work,config: Cost.Configuration) -> Movement:
172	  Work{positions,inner,rows,owners,geometry,activation} = work
173	  movement_with(~Geometry,~fork_window,depth,offset,span,window,geometry,work,config)
174	
175	def workspace(+depth: U32,+work: Work,cost: Movement) -> Nat:
176	  Work{+positions,inner,rows,owners,geometry,activation} = work
177	  Movement{packed,leaves,copies,span,total,minimum} = cost
178	  Nat.add(weight_bytes(work),Nat.add(packed,Nat.add(copies,
179	    Bool.pick(Nat,U32.is_zero(depth),0n,leaves))))
180	
181	# Estimate the partition on workers by (work - span) / workers + span, Graham's
182	# bound for greedy list scheduling. BendRT does not steal work: it deals forked
183	# tasks across rings that workers claim each turn, so this is an estimate, not
184	# a guarantee. One worker gives the total work, so a split only adds its
185	# copies; more workers shorten the span but pay copies.
186	def scheduled(+workers: U32,+span: F32,+total: F32) -> F32:
187	  ((total - span) / U32.to_f32(workers) + span : F32)
188	
189	def estimate(+depth: U32,+work: Work,config: Cost.Configuration,cost: Movement,+workers: U32) -> F32:
190	  Work{+positions,+inner,+rows,owners,geometry,activation} = work
191	  Cost.Configuration{mac,+byte_ns,operation_ns,grain,wake} = config
192	  Movement{packed,leaves,copies,+span,+total,minimum} = cost
193	  +reorder = reorder_ns(work,config)
194	  # A depth-zero blocked leaf already owns the exact logical output.
195	  +gather = Bool.pick(F32,U32.is_zero(depth),0.0,
196	    ((U32.to_f32(rows) * U32.to_f32(positions) * 4.0 + payload_float(Nat.mul(4n,capacity((rows * positions : U32))))) * byte_ns : F32))
197	  (reorder + scheduled(workers,span,total) + gather : F32)
198	
199	def fits(budget: Maybe<&2,Nat>,bytes: Nat) -> Bool:
200	  match budget:
201	    case None{}: True{}
202	    case Some{limit}: Nat.is_le(bytes,limit)
203	
204	def choose(fits: Bool,depth: U32,bytes: Nat,+time: F32,best: Plan) -> Plan:
205	  Plan{+packed,old_depth,old_bytes,+old_time} = best
206	  Bool.pick(Plan,fits && (Bool.not(packed) || F32.is_lt(time,old_time)),Plan{True{},depth,bytes,time},Plan{packed,old_depth,old_bytes,old_time})
207	
208	def candidate(+depth: U32,+work: Work,+config: Cost.Configuration,+budget: Maybe<&2,Nat>,+workers: U32,+cost: Movement,best: Plan) -> Plan:
209	  Cost.Configuration{mac,byte_ns,operation_ns,grain,wake} = config
210	  Movement{packed,leaves,copies,span,total,minimum} = cost
211	  choose((U32.is_zero(depth) || F32.is_le(grain,minimum)) && fits(budget,workspace(depth,work,cost)),
212	    depth,workspace(depth,work,cost),estimate(depth,work,config,cost,workers),best)
213	
214	def candidates_with(~Context: Data,~fork: @+offset: U32 -> @+half: U32 -> @+window: InputRows -> @+geometry: Context -> ForkInput,
215	  remaining: Nat,allowed: Bool,+depth: U32,+tasks: U32,+ceiling: U32,+workers: U32,+span: U32,+geometry: Context,+work: Work,+config: Cost.Configuration,
216	  +budget: Maybe<&2,Nat>,best: Plan) -> Plan:
217	  match remaining allowed:
218	    case 0n allowed: best
219	    case 1n+rest False{}: best
220	    case 1n+rest True{}:
221	      +next_tasks = (tasks * 2 : U32)
222	      candidates_with(~Context,~fork,rest,U32.is_le(tasks,(ceiling / 2 : U32)) && U32.is_le(1,(span / 8 / next_tasks : U32)),
223	        (depth + 1 : U32),next_tasks,ceiling,workers,span,geometry,work,config,budget,
224	        candidate(depth,work,config,budget,workers,movement_with(~Context,~fork,U32.to_nat(depth),0,span,root_rows(work),geometry,work,config),best))
225	
226	def candidates(remaining: Nat,allowed: Bool,depth: U32,tasks: U32,ceiling: U32,workers: U32,span: U32,+work: Work,config: Cost.Configuration,
227	  budget: Maybe<&2,Nat>,best: Plan) -> Plan:
228	  Work{positions,inner,rows,owners,geometry,activation} = work
229	  candidates_with(~Geometry,~fork_window,remaining,allowed,depth,tasks,ceiling,workers,span,geometry,work,config,budget,best)
230	
231	# workers is the runtime's pool size. As in XNNPACK, several workers take up to
232	# five tasks each, so idle workers pick up the remaining ones; one worker never
233	# splits. Every leaf keeps at least one block and one grain of work, and the
234	# scheduling bound must strictly improve.
235	def select(+work: Work,workers: U32,+budget: Maybe<&2,Nat>,+config: Cost.Configuration) -> Plan:
236	  Work{positions,inner,rows,owners,geometry,activation} = work
237	  +workers = U32.max(1,U32.min(workers,858993459))
238	  candidates(32n,True{},0,1,Bool.pick(U32,U32.is_eq(workers,1),1,(workers * 5 : U32)),workers,KernelWindow.padded(positions),work,config,budget,
239	    Plan{False{},0,0n,0.0})