kernel_band_planning.bend
stelliferous@0.0.3.0 · 185 lines · 0 definitions
1# Row-layout costs follow the actual halo ranges and allocation capacities.
2# Preparation is serial. Fork copies weights and bias, never input cores.
3import Base
4import stelliferous@0.0.3.0/kernel_cost.bend as Cost
5import stelliferous@0.0.3.0/kernel_planning.bend as KernelPlanning
6import stelliferous@0.0.3.0/tensor_band.bend as Bands
7import stelliferous@0.0.3.0/tensor_band_window.bend as BandWindow
8import stelliferous@0.0.3.0/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.
12type Movement is Data:
13 Movement{work: F32,span: F32,preparation: F32,minimum: F32,leaves: U32}
14
15type Plan is Data:
16 Plan{depth: U32,parallel: Bool,estimated_ns: F32}
17
18def 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
21def 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.
26def 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
29def 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
36def 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
42def 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
46def 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.
50type Priced is Data:
51 Priced{workers: U32,work: KernelPlanning.Work,window: BandWindow.Geometry,config: Cost.Configuration}
52
53def 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
59def 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
63def 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
66def 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
70def 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.
75type Moved is Data:
76 Moved{geometry: Bands.Geometry,byte_ns: F32}
77
78def 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
82def add_transfer(first: F32,second: F32) -> F32:
83 (first + second : F32)
84
85def 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.
89def wake_cost(config: Cost.Configuration) -> F32:
90 Cost.Configuration{mac,byte_ns,operation_ns,grain,wake} = config
91 (2.0 * wake : F32)
92
93def 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
98def 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
103def 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
113def 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.
124def 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
131type LayoutPlan is Data:
132 LayoutPlan{layout: Bands.Plan,parallel: Bool,estimated_ns: F32,repartition_ns: F32}
133
134def 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
142def 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
145def 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
150def 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
155def 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
158def 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
168def 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.
175def 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.
182def 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)