1	# Row-layout costs follow the actual halo ranges and allocation capacities.
2	# Preparation is serial. Fork copies weights and bias, never input cores.
3	import Base
4	import ./kernel_cost.bend as Cost
5	import ./kernel_planning.bend as KernelPlanning
6	import ./tensor_band.bend as Bands
7	import ./tensor_band_window.bend as BandWindow
8	import ./kernel_window.bend as KernelWindow
9	
10	# The whole batch is one parallel region: span is its costliest leaf path,
11	# leaves its leaves, across every image.
12	type Movement is Data:
13	  Movement{work: F32,span: F32,preparation: F32,minimum: F32,leaves: U32}
14	
15	type Plan is Data:
16	  Plan{depth: U32,parallel: Bool,estimated_ns: F32}
17	
18	def transfer(+count: U32,byte_ns: F32) -> F32:
19	  ((U32.to_f32(count) * 4.0 + KernelPlanning.payload_float(Nat.mul(4n,KernelPlanning.capacity(count)))) * byte_ns : F32)
20	
21	def leaf_cost(preparation: F32,+tail: F32,cost: KernelPlanning.Movement) -> Movement:
22	  KernelPlanning.Movement{packed,output,copies,span,total,minimum} = cost
23	  Movement{(total + tail : F32),(span + tail : F32),preparation,minimum,1}
24	
25	# BandWindow.assemble reuses the core exactly when both row ranges agree.
26	def assembly(channels: U32,width: U32,first: U32,rows: U32,window_first: U32,+window_rows: U32,byte_ns: F32) -> F32:
27	  Bool.pick(F32,U32.is_eq(first,window_first) && U32.is_eq(rows,window_rows),0.0,transfer((channels * window_rows * width : U32),byte_ns))
28	
29	def leaf_window(+work: KernelPlanning.Work,first: U32,count: U32,window: BandWindow.RowWindow,+config: Cost.Configuration) -> Movement:
30	  match work window config:
31	    case KernelPlanning.Work{positions,inner,rows,owners,KernelPlanning.Geometry{+channels,h,+width,k,s,p,+ow},activation} BandWindow.RowWindow{BandWindow.Rows{out_first,+out_count},BandWindow.Rows{af,+ac},BandWindow.Rows{bf,+bc},BandWindow.Rows{wf,+wc}} Cost.Configuration{mac,+byte_ns,operation_ns,grain,wake}:
32	      leaf_cost((transfer((channels * ac * width : U32),byte_ns) + transfer((channels * bc * width : U32),byte_ns) : F32),
33	        assembly(channels,width,first,count,wf,wc,byte_ns),
34	        KernelPlanning.leaf_cost(KernelWindow.padded((out_count * ow : U32)),(out_count * ow : U32),work,config))
35	
36	def merge(clone: F32,left: Movement,right: Movement) -> Movement:
37	  Movement{lw,ls,lp,lmin,ln} = left
38	  Movement{rw,rs,rp,rmin,rn} = right
39	  +clone = clone
40	  Movement{(clone + lw + rw : F32),(clone + F32.max(ls,rs) : F32),(lp + rp : F32),F32.min(lmin,rmin),(ln + rn : U32)}
41	
42	def clone_cost(work: KernelPlanning.Work,config: Cost.Configuration) -> F32:
43	  Cost.Configuration{mac,byte_ns,operation_ns,grain,wake} = config
44	  (KernelPlanning.payload_float(KernelPlanning.fork_bytes(work,0n)) * byte_ns : F32)
45	
46	def reorder_cost(work: KernelPlanning.Work,config: Cost.Configuration) -> F32:
47	  KernelPlanning.reorder_ns(work,config)
48	
49	# What the movement of a layout is priced with.
50	type Priced is Data:
51	  Priced{workers: U32,work: KernelPlanning.Work,window: BandWindow.Geometry,config: Cost.Configuration}
52	
53	def leaf_movement(context: Priced,first: U32,count: U32) -> Movement:
54	  Priced{workers,work,+window,config} = context
55	  +first = first
56	  +count = count
57	  leaf_window(work,first,count,BandWindow.window(window,first,count),config)
58	
59	def split_movement(context: Priced,left: Bands.Plan,right: Bands.Plan,first: Movement,second: Movement) -> Movement:
60	  Priced{workers,+work,window,+config} = context
61	  merge(clone_cost(work,config),first,second)
62	
63	def movement(layout: Bands.Plan,+workers: U32,+work: KernelPlanning.Work,+window: BandWindow.Geometry,+config: Cost.Configuration) -> Movement:
64	  Bands.plan_fold(~Movement,~Priced,~leaf_movement,~split_movement,layout,Priced{workers,work,window,config})
65	
66	def estimate(parallel: Bool,+workers: U32,initial_copy: F32,cost: Movement) -> F32:
67	  Movement{+work,+span,preparation,minimum,+leaves} = cost
68	  (initial_copy + preparation + Bool.pick(F32,parallel,Cost.region_ns(workers,leaves,span),work) : F32)
69	
70	def choose(allowed: Bool,depth: U32,parallel: Bool,+time: F32,best: Plan) -> Plan:
71	  Plan{old_depth,old_parallel,+old_time} = best
72	  Bool.pick(Plan,allowed && F32.is_lt(time,old_time),Plan{depth,parallel,time},Plan{old_depth,old_parallel,old_time})
73	
74	# What a transfer of bands is priced with.
75	type Moved is Data:
76	  Moved{geometry: Bands.Geometry,byte_ns: F32}
77	
78	def leaf_transfer(context: Moved,first: U32,rows: U32) -> F32:
79	  Moved{geometry,byte_ns} = context
80	  transfer(Bands.local_count(geometry,rows),byte_ns)
81	
82	def add_transfer(first: F32,second: F32) -> F32:
83	  (first + second : F32)
84	
85	def transfer_layout(layout: Bands.Plan,+geometry: Bands.Geometry,+byte_ns: F32) -> F32:
86	  Bands.plan_fold(~F32,~Moved,~leaf_transfer,~Bands.plain(~F32,~Moved,~add_transfer),layout,Moved{geometry,byte_ns})
87	
88	# The two wake rounds of the one parallel region.
89	def wake_cost(config: Cost.Configuration) -> F32:
90	  Cost.Configuration{mac,byte_ns,operation_ns,grain,wake} = config
91	  (2.0 * wake : F32)
92	
93	def candidate(+depth: U32,+workers: U32,+initial_copy: F32,+wake: F32,+cost: Movement,config: Cost.Configuration,best: Plan) -> Plan:
94	  match cost config:
95	    case Movement{work,span,preparation,minimum,leaves} Cost.Configuration{mac,byte_ns,operation_ns,grain,round}:
96	      choose(F32.is_le(grain,work),depth,True{},(estimate(True{},workers,initial_copy,cost) + wake : F32),best)
97	
98	def initial_candidate(depth: U32,+layout: Bands.Plan,+workers: U32,+work: KernelPlanning.Work,+window: BandWindow.Geometry,+config: Cost.Configuration,best: Plan) -> Plan:
99	  Cost.Configuration{mac,byte_ns,operation_ns,grain,wake} = config
100	  candidate(depth,workers,(transfer_layout(layout,BandWindow.input_geometry(window),byte_ns) + reorder_cost(work,config) : F32),
101	    wake_cost(config),movement(layout,workers,work,window,config),config,best)
102	
103	def candidates(remaining: Nat,allowed: Bool,+depth: U32,+tasks: U32,+ceiling: U32,+workers: U32,+batch: U32,+height: U32,
104	  +work: KernelPlanning.Work,+window: BandWindow.Geometry,+config: Cost.Configuration,best: Plan) -> Plan:
105	  match remaining allowed:
106	    case 0n allowed: best
107	    case 1n+rest False{}: best
108	    case 1n+rest True{}:
109	      +next_tasks = (tasks * 2 : U32)
110	      candidates(rest,U32.is_le(next_tasks,ceiling) && U32.is_le(next_tasks,height),(depth + 1 : U32),next_tasks,ceiling,workers,batch,height,
111	        work,window,config,initial_candidate(depth,Bands.image_plan(U32.to_nat(batch),0,U32.to_nat(depth),height),workers,work,window,config,best))
112	
113	def unsplit_baseline(+work: KernelPlanning.Work,window: BandWindow.Geometry,+config: Cost.Configuration) -> F32:
114	  match work window config:
115	    case KernelPlanning.Work{+positions,inner,+rows,owners,geometry,activation} BandWindow.Geometry{Bands.Geometry{+batch,channels,height,width},kernel,stride,padding,oh,ow} Cost.Configuration{mac,+byte_ns,operation_ns,grain,wake}:
116	      +one = KernelPlanning.estimate(0,work,config,KernelPlanning.movement(0n,0,KernelWindow.padded(positions),KernelPlanning.root_rows(work),work,config),1)
117	      +output_size = (rows * positions : U32)
118	      +images = U32.to_f32(batch)
119	      Bool.pick(F32,U32.is_le(batch,1),one,
120	        (reorder_cost(work,config) + images * (one - reorder_cost(work,config) + transfer((channels * height * width : U32),byte_ns) + U32.to_f32(output_size) * 4.0 * byte_ns) +
121	          KernelPlanning.payload_float(Nat.mul(4n,KernelPlanning.capacity((batch * output_size : U32)))) * byte_ns + (images - 1.0) * clone_cost(work,config) : F32))
122	
123	# Grain bounds the operation's work; the two wake rounds are charged separately.
124	def select(+work: KernelPlanning.Work,+window: BandWindow.Geometry,workers: U32,+config: Cost.Configuration) -> Plan:
125	  match work window:
126	    case KernelPlanning.Work{positions,inner,rows,owners,KernelPlanning.Geometry{channels,+height,width,k,s,p,ow},activation} BandWindow.Geometry{Bands.Geometry{batch,ci,h,w},kernel,stride,padding,oh,owidth}:
127	      +workers = U32.max(1,U32.min(workers,858993459))
128	      candidates(32n,U32.is_lt(1,workers) && (U32.is_lt(1,batch) || U32.is_le(2,height)),0,1,(workers * 5 : U32),workers,batch,height,work,window,config,
129	        Plan{0,False{},unsplit_baseline(work,window,config)})
130	
131	type LayoutPlan is Data:
132	  LayoutPlan{layout: Bands.Plan,parallel: Bool,estimated_ns: F32,repartition_ns: F32}
133	
134	def retained_parallel(window: BandWindow.Geometry,workers: U32,config: Cost.Configuration,cost: Movement) -> Bool:
135	  BandWindow.Geometry{Bands.Geometry{+batch,channels,height,width},kernel,stride,padding,oh,ow} = window
136	  Cost.Configuration{mac,byte_ns,operation_ns,grain,wake} = config
137	  Movement{+total,+span,preparation,minimum,+leaves} = cost
138	  +workers = U32.max(1,U32.min(workers,858993459))
139	  U32.is_lt(1,workers) && Nat.is_le(U32.to_nat(leaves),Nat.mul(U32.to_nat(batch),U32.to_nat((workers * 5 : U32)))) &&
140	    F32.is_le(grain,total) && F32.is_lt((Cost.region_ns(workers,leaves,span) + 2.0 * wake : F32),total)
141	
142	def layout_cost(layout: Bands.Plan,+parallel: Bool,workers: U32,+repartition: F32,wake: F32,reorder: F32,cost: Movement) -> LayoutPlan:
143	  LayoutPlan{layout,parallel,(reorder + estimate(parallel,workers,repartition,cost) + Bool.pick(F32,parallel,wake,0.0) : F32),repartition}
144	
145	def evaluated(+layout: Bands.Plan,+previous: Bands.Plan,+work: KernelPlanning.Work,+window: BandWindow.Geometry,+workers: U32,+config: Cost.Configuration,+cost: Movement) -> LayoutPlan:
146	  Cost.Configuration{mac,byte_ns,operation_ns,grain,wake} = config
147	  layout_cost(layout,retained_parallel(window,workers,config,cost),workers,
148	    Bool.pick(F32,Bands.same_plan(layout,previous),0.0,transfer_layout(layout,BandWindow.input_geometry(window),byte_ns)),wake_cost(config),reorder_cost(work,config),cost)
149	
150	def prefer(candidate: LayoutPlan,best: LayoutPlan) -> LayoutPlan:
151	  match candidate best:
152	    case LayoutPlan{layout,parallel,+time,copies} LayoutPlan{old,old_parallel,+old_time,old_copies}:
153	      Bool.pick(LayoutPlan,F32.is_lt(time,old_time),LayoutPlan{layout,parallel,time,copies},LayoutPlan{old,old_parallel,old_time,old_copies})
154	
155	def layout_candidate(+layout: Bands.Plan,previous: Bands.Plan,+work: KernelPlanning.Work,+window: BandWindow.Geometry,+workers: U32,+config: Cost.Configuration,best: LayoutPlan) -> LayoutPlan:
156	  prefer(evaluated(layout,previous,work,window,workers,config,movement(layout,workers,work,window,config)),best)
157	
158	def layouts(remaining: Nat,allowed: Bool,+depth: U32,+tasks: U32,+ceiling: U32,+height: U32,+batch: U32,+previous: Bands.Plan,
159	  +work: KernelPlanning.Work,+window: BandWindow.Geometry,+workers: U32,+config: Cost.Configuration,best: LayoutPlan) -> LayoutPlan:
160	  match remaining allowed:
161	    case 0n allowed: best
162	    case 1n+rest False{}: best
163	    case 1n+rest True{}:
164	      +next_tasks = (tasks * 2 : U32)
165	      layouts(rest,U32.is_le(next_tasks,ceiling) && U32.is_le(next_tasks,height),(depth + 1 : U32),next_tasks,ceiling,height,batch,previous,work,window,workers,config,
166	        layout_candidate(Bands.image_plan(U32.to_nat(batch),0,U32.to_nat(depth),height),previous,work,window,workers,config,best))
167	
168	def replan(+previous: Bands.Plan,+work: KernelPlanning.Work,+window: BandWindow.Geometry,workers: U32,+config: Cost.Configuration) -> LayoutPlan:
169	  BandWindow.Geometry{Bands.Geometry{batch,channels,height,width},kernel,stride,padding,oh,ow} = window
170	  +workers = U32.max(1,U32.min(workers,858993459))
171	  layouts(32n,True{},0,1,Bool.pick(U32,U32.is_eq(workers,1),1,(workers * 5 : U32)),height,batch,previous,work,window,workers,config,
172	    evaluated(previous,previous,work,window,workers,config,movement(previous,workers,work,window,config)))
173	
174	# The first split of every image as a layout plan: nothing to repartition.
175	def first_split(plan: Plan,+window: BandWindow.Geometry) -> LayoutPlan:
176	  Plan{depth,parallel,time} = plan
177	  BandWindow.Geometry{Bands.Geometry{batch,channels,+height,width},kernel,stride,padding,oh,ow} = window
178	  LayoutPlan{Bands.image_plan(U32.to_nat(batch),0,U32.to_nat(depth),height),parallel,time,0.0}
179	
180	# The layout a window operation runs in: chosen again from the input's plan
181	# when the input holds the window's geometry, otherwise the first split.
182	def layout_for(same: Bool,+previous: Bands.Plan,+work: KernelPlanning.Work,+window: BandWindow.Geometry,workers: U32,+config: Cost.Configuration) -> LayoutPlan:
183	  match same:
184	    case True{}: replan(previous,work,window,workers,config)
185	    case False{}: first_split(select(work,window,workers,config),window)