kernel_band_planning.bend

stelliferous@0.0.2.0 · 171 lines · 30 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.2.0/kernel_cost.bend as Cost
5import stelliferous@0.0.2.0/kernel_planning.bend as KernelPlanning
6import stelliferous@0.0.2.0/tensor_band.bend as Bands
7import stelliferous@0.0.2.0/tensor_band_window.bend as BandWindow
8import stelliferous@0.0.2.0/kernel_window.bend as KernelWindow
9
10# Images run one after another, each as its own parallel region (the runtime
11# forks only within an image). span and open describe the current image's
12# region: its costliest leaf path and its leaves. settled is the time of the
13# image regions already finished.
14type Movement is Data:
15  Movement{work: F32,span: F32,open: U32,settled: F32,preparation: F32,minimum: F32,leaves: U32}
16
17type Plan is Data:
18  Plan{depth: U32,parallel: Bool,estimated_ns: F32}
19
20def transfer(+count: U32,byte_ns: F32) -> F32:
21  ((U32.to_f32(count) * 4.0 + KernelPlanning.payload_float(Nat.mul(4n,KernelPlanning.capacity(count)))) * byte_ns : F32)
22
23def leaf_cost(preparation: F32,+tail: F32,cost: KernelPlanning.Movement) -> Movement:
24  KernelPlanning.Movement{packed,output,copies,span,total,minimum} = cost
25  Movement{(total + tail : F32),(span + tail : F32),1,0.0,preparation,minimum,1}
26
27# BandWindow.assemble reuses the core exactly when both row ranges agree.
28def assembly(channels: U32,width: U32,first: U32,rows: U32,window_first: U32,+window_rows: U32,byte_ns: F32) -> F32:
29  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))
30
31def leaf_window(+work: KernelPlanning.Work,first: U32,count: U32,window: BandWindow.RowWindow,+config: Cost.Configuration) -> Movement:
32  match work window config:
33    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}:
34      leaf_cost((transfer((channels * ac * width : U32),byte_ns) + transfer((channels * bc * width : U32),byte_ns) : F32),
35        assembly(channels,width,first,count,wf,wc,byte_ns),
36        KernelPlanning.leaf_cost(KernelWindow.padded((out_count * ow : U32)),(out_count * ow : U32),work,config))
37
38def merge(clone: F32,left: Movement,right: Movement) -> Movement:
39  Movement{lw,ls,lo,lsettled,lp,lmin,ln} = left
40  Movement{rw,rs,ro,rsettled,rp,rmin,rn} = right
41  +clone = clone
42  Movement{(clone + lw + rw : F32),(clone + F32.max(ls,rs) : F32),(lo + ro : U32),(lsettled + rsettled : F32),(lp + rp : F32),F32.min(lmin,rmin),(ln + rn : U32)}
43
44# A split between images finishes both sides' regions in turn.
45def merge_images(+workers: U32,clone: F32,left: Movement,right: Movement) -> Movement:
46  Movement{lw,+ls,+lo,lsettled,lp,lmin,ln} = left
47  Movement{rw,+rs,+ro,rsettled,rp,rmin,rn} = right
48  +clone = clone
49  Movement{(clone + lw + rw : F32),0.0,0,(clone + lsettled + rsettled + Cost.region_ns(workers,lo,ls) + Cost.region_ns(workers,ro,rs) : F32),
50    (lp + rp : F32),F32.min(lmin,rmin),(ln + rn : U32)}
51
52def merge_layout(same_image: Bool,+workers: U32,clone: F32,left: Movement,right: Movement) -> Movement:
53  match same_image:
54    case True{}: merge(clone,left,right)
55    case False{}: merge_images(workers,clone,left,right)
56
57def clone_cost(work: KernelPlanning.Work,config: Cost.Configuration) -> F32:
58  Cost.Configuration{mac,byte_ns,operation_ns,grain,wake} = config
59  (KernelPlanning.payload_float(KernelPlanning.fork_bytes(work,0n)) * byte_ns : F32)
60
61def reorder_cost(work: KernelPlanning.Work,config: Cost.Configuration) -> F32:
62  KernelPlanning.reorder_ns(work,config)
63
64def movement(layout: Bands.Plan,+workers: U32,+work: KernelPlanning.Work,+window: BandWindow.Geometry,+config: Cost.Configuration) -> Movement:
65  match layout:
66    case Bands.Rows{+first,+count}: leaf_window(work,first,count,BandWindow.window(window,first,count),config)
67    case Bands.Split{+left,+right}:
68      merge_layout(Bands.same_image(BandWindow.input_geometry(window),Bands.Split{left,right}),workers,clone_cost(work,config),
69        movement(left,workers,work,window,config),movement(right,workers,work,window,config))
70
71def parallel_ns(+workers: U32,+span: F32,+open: U32,+settled: F32) -> F32:
72  (settled + Cost.region_ns(workers,open,span) : F32)
73
74def estimate(parallel: Bool,+workers: U32,initial_copy: F32,cost: Movement) -> F32:
75  Movement{+work,+span,+open,+settled,preparation,minimum,leaves} = cost
76  (initial_copy + preparation + Bool.pick(F32,parallel,parallel_ns(workers,span,open,settled),work) : F32)
77
78def choose(allowed: Bool,depth: U32,parallel: Bool,+time: F32,best: Plan) -> Plan:
79  Plan{old_depth,old_parallel,+old_time} = best
80  Bool.pick(Plan,allowed && F32.is_lt(time,old_time),Plan{depth,parallel,time},Plan{old_depth,old_parallel,old_time})
81
82def transfer_layout(layout: Bands.Plan,+geometry: Bands.Geometry,+byte_ns: F32) -> F32:
83  match layout:
84    case Bands.Rows{first,rows}: transfer(Bands.local_count(geometry,rows),byte_ns)
85    case Bands.Split{left,right}: (transfer_layout(left,geometry,byte_ns) + transfer_layout(right,geometry,byte_ns) : F32)
86
87def wake_cost(window: BandWindow.Geometry,config: Cost.Configuration) -> F32:
88  BandWindow.Geometry{Bands.Geometry{batch,channels,height,width},kernel,stride,padding,oh,ow} = window
89  Cost.Configuration{mac,byte_ns,operation_ns,grain,wake} = config
90  (2.0 * U32.to_f32(batch) * wake : F32)
91
92def candidate(+depth: U32,+workers: U32,+initial_copy: F32,+wake: F32,+cost: Movement,config: Cost.Configuration,best: Plan) -> Plan:
93  match cost config:
94    case Movement{work,span,open,settled,preparation,minimum,leaves} Cost.Configuration{mac,byte_ns,operation_ns,grain,round}:
95      choose(F32.is_le(grain,work),depth,True{},(estimate(True{},workers,initial_copy,cost) + wake : F32),best)
96
97def initial_candidate(depth: U32,+layout: Bands.Plan,+workers: U32,+work: KernelPlanning.Work,+window: BandWindow.Geometry,+config: Cost.Configuration,best: Plan) -> Plan:
98  Cost.Configuration{mac,byte_ns,operation_ns,grain,wake} = config
99  candidate(depth,workers,(transfer_layout(layout,BandWindow.input_geometry(window),byte_ns) + reorder_cost(work,config) : F32),
100    wake_cost(window,config),movement(layout,workers,work,window,config),config,best)
101
102def candidates(remaining: Nat,allowed: Bool,+depth: U32,+tasks: U32,+ceiling: U32,+workers: U32,+batch: U32,+height: U32,
103  +work: KernelPlanning.Work,+window: BandWindow.Geometry,+config: Cost.Configuration,best: Plan) -> Plan:
104  match remaining allowed:
105    case 0n allowed: best
106    case 1n+rest False{}: best
107    case 1n+rest True{}:
108      +next_tasks = (tasks * 2 : U32)
109      candidates(rest,U32.is_le(next_tasks,ceiling) && U32.is_le(next_tasks,height),(depth + 1 : U32),next_tasks,ceiling,workers,batch,height,
110        work,window,config,initial_candidate(depth,Bands.image_plan(U32.to_nat(batch),0,U32.to_nat(depth),height),workers,work,window,config,best))
111
112def dense_baseline(+work: KernelPlanning.Work,window: BandWindow.Geometry,+config: Cost.Configuration) -> F32:
113  match work window config:
114    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}:
115      +one = KernelPlanning.estimate(0,work,config,KernelPlanning.movement(0n,0,KernelWindow.padded(positions),KernelPlanning.root_rows(work),work,config),1)
116      +output_size = (rows * positions : U32)
117      +images = U32.to_f32(batch)
118      Bool.pick(F32,U32.is_le(batch,1),one,
119        (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) +
120          KernelPlanning.payload_float(Nat.mul(4n,KernelPlanning.capacity((batch * output_size : U32)))) * byte_ns + (images - 1.0) * clone_cost(work,config) : F32))
121
122# Grain bounds the operation's work; the two wake rounds are charged separately.
123def select(+work: KernelPlanning.Work,+window: BandWindow.Geometry,workers: U32,+config: Cost.Configuration) -> Plan:
124  match work window:
125    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}:
126      +workers = U32.max(1,U32.min(workers,858993459))
127      candidates(31n,U32.is_lt(1,workers) && U32.is_le(2,height),1,2,(workers * 5 : U32),workers,batch,height,work,window,config,
128        Plan{0,False{},dense_baseline(work,window,config)})
129
130type LayoutPlan is Data:
131  LayoutPlan{layout: Bands.Plan,parallel: Bool,estimated_ns: F32,repartition_ns: F32}
132
133def retained_parallel(window: BandWindow.Geometry,workers: U32,config: Cost.Configuration,cost: Movement) -> Bool:
134  BandWindow.Geometry{Bands.Geometry{+batch,channels,height,width},kernel,stride,padding,oh,ow} = window
135  Cost.Configuration{mac,byte_ns,operation_ns,grain,wake} = config
136  Movement{+total,+span,+open,+settled,preparation,minimum,+leaves} = cost
137  +workers = U32.max(1,U32.min(workers,858993459))
138  U32.is_lt(1,workers) && Nat.is_le(U32.to_nat(leaves),Nat.mul(U32.to_nat(batch),U32.to_nat((workers * 5 : U32)))) &&
139    F32.is_le(grain,total) && F32.is_lt((parallel_ns(workers,span,open,settled) + 2.0 * U32.to_f32(batch) * wake : F32),total)
140
141def layout_cost(layout: Bands.Plan,+parallel: Bool,workers: U32,+repartition: F32,wake: F32,reorder: F32,cost: Movement) -> LayoutPlan:
142  LayoutPlan{layout,parallel,(reorder + estimate(parallel,workers,repartition,cost) + Bool.pick(F32,parallel,wake,0.0) : F32),repartition}
143
144def evaluated(+layout: Bands.Plan,+previous: Bands.Plan,+work: KernelPlanning.Work,+window: BandWindow.Geometry,+workers: U32,+config: Cost.Configuration,+cost: Movement) -> LayoutPlan:
145  Cost.Configuration{mac,byte_ns,operation_ns,grain,wake} = config
146  layout_cost(layout,retained_parallel(window,workers,config,cost),workers,
147    Bool.pick(F32,Bands.same_plan(layout,previous),0.0,transfer_layout(layout,BandWindow.input_geometry(window),byte_ns)),wake_cost(window,config),reorder_cost(work,config),cost)
148
149def prefer(candidate: LayoutPlan,best: LayoutPlan) -> LayoutPlan:
150  match candidate best:
151    case LayoutPlan{layout,parallel,+time,copies} LayoutPlan{old,old_parallel,+old_time,old_copies}:
152      Bool.pick(LayoutPlan,F32.is_lt(time,old_time),LayoutPlan{layout,parallel,time,copies},LayoutPlan{old,old_parallel,old_time,old_copies})
153
154def layout_candidate(+layout: Bands.Plan,previous: Bands.Plan,+work: KernelPlanning.Work,+window: BandWindow.Geometry,+workers: U32,+config: Cost.Configuration,best: LayoutPlan) -> LayoutPlan:
155  prefer(evaluated(layout,previous,work,window,workers,config,movement(layout,workers,work,window,config)),best)
156
157def layouts(remaining: Nat,allowed: Bool,+depth: U32,+tasks: U32,+ceiling: U32,+height: U32,+batch: U32,+previous: Bands.Plan,
158  +work: KernelPlanning.Work,+window: BandWindow.Geometry,+workers: U32,+config: Cost.Configuration,best: LayoutPlan) -> LayoutPlan:
159  match remaining allowed:
160    case 0n allowed: best
161    case 1n+rest False{}: best
162    case 1n+rest True{}:
163      +next_tasks = (tasks * 2 : U32)
164      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,
165        layout_candidate(Bands.image_plan(U32.to_nat(batch),0,U32.to_nat(depth),height),previous,work,window,workers,config,best))
166
167def replan(+previous: Bands.Plan,+work: KernelPlanning.Work,+window: BandWindow.Geometry,workers: U32,+config: Cost.Configuration) -> LayoutPlan:
168  BandWindow.Geometry{Bands.Geometry{batch,channels,height,width},kernel,stride,padding,oh,ow} = window
169  +workers = U32.max(1,U32.min(workers,858993459))
170  layouts(32n,True{},0,1,Bool.pick(U32,U32.is_eq(workers,1),1,(workers * 5 : U32)),height,batch,previous,work,window,workers,config,
171    evaluated(previous,previous,work,window,workers,config,movement(previous,workers,work,window,config)))