tensor_band_schedule.bend
stelliferous@0.0.3.0 · 124 lines · 0 definitions
1# Conservative memory-work estimates for elementwise and layout operations.
2# Small operations retain their owners and run sequentially. Multi-image trees
3# remain sequential here; window operations have explicit per-image dispatch.
4import Base
5import stelliferous@0.0.3.0/tensor_band.bend as Bands
6import stelliferous@0.0.3.0/storage_buffer.bend as Storage
7import stelliferous@0.0.3.0/kernel_cost.bend as Cost
8import stelliferous@0.0.3.0/default_kernel_cost.bend as DefaultCost
9
10type Size is Data:
11 Size{leaves: U32,minimum_rows: U32,maximum_rows: U32}
12
13def merge(left: Size,right: Size) -> Size:
14 Size{ln,lmin,lmax} = left
15 Size{rn,rmin,rmax} = right
16 Size{(ln + rn : U32),U32.min(lmin,rmin),U32.max(lmax,rmax)}
17
18def leaf_size(context: Unit,first: U32,rows: U32) -> Size:
19 +rows = rows
20 Size{1,rows,rows}
21
22def size(layout: Bands.Plan) -> Size:
23 Bands.plan_fold(~Size,~Unit,~leaf_size,~Bands.plain(~Size,~Unit,~merge),layout,Unit{})
24
25# The whole operation, not each leaf, must reach the grain: small leaves only
26# lengthen the runtime's rounds, which the region estimate already counts.
27def sized(geometry: Bands.Geometry,+workers: U32,element_ns: F32,config: Cost.Configuration,size: Size) -> Bool:
28 Bands.Geometry{+batch,channels,height,width} = geometry
29 Cost.Configuration{mac,byte_ns,operation_ns,grain,wake} = config
30 Size{+leaves,rows,maximum} = size
31 +row_ns = (U32.to_f32(channels) * U32.to_f32(width) * element_ns : F32)
32 +total = (U32.to_f32(batch) * U32.to_f32(height) * row_ns : F32)
33 U32.is_lt(1,workers) && U32.is_lt(1,leaves) && Nat.is_le(U32.to_nat(leaves),Nat.mul(U32.to_nat(batch),U32.to_nat((U32.min(workers,858993459) * 5 : U32)))) &&
34 F32.is_le(grain,total) &&
35 F32.is_lt((Cost.region_ns(workers,leaves,(U32.to_f32(maximum) * row_ns : F32)) + 2.0 * wake : F32),total)
36
37def parallel(geometry: Bands.Geometry,workers: U32,bytes_per_element: F32,operations: F32,layout: Bands.Plan) -> Bool:
38 +config = DefaultCost.configuration()
39 sized(geometry,workers,Cost.element_ns(bytes_per_element,operations,config),config,size(layout))
40
41# For equal-shaped operands the logical read/write volume of repartition is
42# identical. The target owners' rounded capacities decide the remaining cost.
43def leaf_elements(geometry: Bands.Geometry,first: U32,rows: U32) -> Nat:
44 U32.to_nat(U32.shln(1,Storage.depth_for(Bands.local_count(geometry,rows))))
45
46def allocation_elements(layout: Bands.Plan,+geometry: Bands.Geometry) -> Nat:
47 Bands.plan_fold(~Nat,~Bands.Geometry,~leaf_elements,~Bands.plain(~Nat,~Bands.Geometry,~Nat.add),layout,geometry)
48
49def alignment(+geometry: Bands.Geometry,+left: Bands.Plan,+right: Bands.Plan) -> Bands.Plan:
50 Bool.pick(Bands.Plan,Bands.same_plan(left,right) || Nat.is_le(allocation_elements(left,geometry),allocation_elements(right,geometry)),left,right)
51
52# The layout an operation runs in: a geometry, a plan of its rows, and whether
53# the plan splits them into more than one band.
54type Target is Data:
55 Target{geometry: Bands.Geometry,layout: Bands.Plan,split: Bool}
56
57def splits(layout: Bands.Plan) -> Bool:
58 match layout:
59 case Bands.Rows{first,count}: False{}
60 case Bands.Split{left,right}: True{}
61
62def joined(same: Bool,+geometry: Bands.Geometry,+layout: Bands.Plan,other_layout: Bands.Plan) -> Target:
63 match same:
64 case True{}: Target{geometry,alignment(geometry,layout,other_layout),True{}}
65 case False{}: Target{geometry,layout,True{}}
66
67def met(split: Bool,other_split: Bool,+geometry: Bands.Geometry,layout: Bands.Plan,other: Bands.Geometry,other_layout: Bands.Plan) -> Target:
68 match split other_split:
69 case True{} True{}: joined(Bands.same_geometry(geometry,other),geometry,layout,other_layout)
70 case True{} False{}: Target{geometry,layout,True{}}
71 case False{} True{}: Target{other,other_layout,True{}}
72 case False{} False{}: Target{geometry,layout,False{}}
73
74# The layout of two operands: one split into bands names it, and of two split
75# operands of one geometry the plan that allocates less. When neither is
76# split, the first names it.
77def meet(left: Target,right: Target) -> Target:
78 Target{geometry,layout,split} = left
79 Target{other,other_layout,other_split} = right
80 met(split,other_split,geometry,layout,other,other_layout)
81
82def plane_chosen(rows: Bool,+planes: U32,+height: U32,+width: U32,target: Target) -> Target:
83 match rows:
84 case True{}: target
85 case False{}: Target{Bands.Geometry{1,planes,height,width},Bands.unsplit(Bands.Geometry{1,planes,height,width}),False{}}
86
87# The layout of an operation on each plane of the last two axes: the tensor's
88# own when its bands split the image rows of its shape, otherwise its planes
89# as the channels of one image in one band, which a whole tensor already is.
90def plane_target(+shape: List<&2,U32>,+planes: U32,+height: U32,+width: U32,target: Target) -> Target:
91 Target{+geometry,layout,split} = target
92 plane_chosen(Bands.image_rows(shape,geometry),planes,height,width,Target{geometry,layout,split})
93
94def mapped_context(~Element: Data,~Context: Data,~mapper: @+count: Nat -> @+position: U32 -> @+context: Context -> Storage.Buffer<Element> -> Storage.Buffer<Element>,context: Context,+geometry: Bands.Geometry,workers: U32,operations: F32,owners: Bands.Bands(~Element) & Bands.Plan) -> Bands.Bands(~Element):
95 (bands,+layout) = owners
96 Bands.map_with(~Element,~Context,~Bands.map_whole_context(~Element,~Context,~mapper),layout,parallel(geometry,workers,8.0,operations,layout),context,bands)
97
98def map_context(~Element: Data,~Context: Data,~mapper: @+count: Nat -> @+position: U32 -> @+context: Context -> Storage.Buffer<Element> -> Storage.Buffer<Element>,context: Context,geometry: Bands.Geometry,workers: U32,operations: F32,bands: Bands.Bands(~Element)) -> Bands.Bands(~Element):
99 mapped_context(~Element,~Context,~mapper,context,geometry,workers,operations,Bands.plan(~Element,bands))
100
101def scalar_planned(~Element: Data,~operation: Element -> Element -> Element,+geometry: Bands.Geometry,workers: U32,operations: F32,value: Element,owners: Bands.Bands(~Element) & Bands.Plan) -> Bands.Bands(~Element):
102 (bands,+layout) = owners
103 Bands.scalar(~Element,~operation,layout,parallel(geometry,workers,8.0,operations,layout),value,bands)
104
105def scalar(~Element: Data,~operation: Element -> Element -> Element,geometry: Bands.Geometry,workers: U32,operations: F32,value: Element,bands: Bands.Bands(~Element)) -> Bands.Bands(~Element):
106 scalar_planned(~Element,~operation,geometry,workers,operations,value,Bands.plan(~Element,bands))
107
108def indexed_planned(~Element: Data,~operation: U32 -> Element -> Element,+geometry: Bands.Geometry,workers: U32,operations: F32,fill: Element,owners: Bands.Bands(~Element) & Bands.Plan) -> Bands.Bands(~Element):
109 (bands,+layout) = owners
110 Bands.apply_indexed(~Element,~operation,layout,parallel(geometry,workers,12.0,operations,layout),geometry,fill,bands)
111
112def indexed(~Element: Data,~operation: U32 -> Element -> Element,geometry: Bands.Geometry,workers: U32,operations: F32,fill: Element,bands: Bands.Bands(~Element)) -> Bands.Bands(~Element):
113 indexed_planned(~Element,~operation,geometry,workers,operations,fill,Bands.plan(~Element,bands))
114
115def transformed(~Element: Data,~Context: Data,~leaf: Context -> U32 -> U32 -> Storage.Buffer<Element> -> Bands.Transformed(~Element),
116 context: Context,+geometry: Bands.Geometry,workers: U32,bytes_per_element: F32,operations: F32,owners: Bands.Bands(~Element) & Bands.Plan) -> Bands.Bands(~Element) & Bool:
117 (bands,+layout) = owners
118 Bands.transform(~Element,~Context,~leaf,layout,parallel(geometry,workers,bytes_per_element,operations,layout),context,bands)
119
120# Every band runs the leaf in the plan it has. bytes_per_element and
121# operations price one input element.
122def transform(~Element: Data,~Context: Data,~leaf: Context -> U32 -> U32 -> Storage.Buffer<Element> -> Bands.Transformed(~Element),
123 context: Context,geometry: Bands.Geometry,workers: U32,bytes_per_element: F32,operations: F32,bands: Bands.Bands(~Element)) -> Bands.Bands(~Element) & Bool:
124 transformed(~Element,~Context,~leaf,context,geometry,workers,bytes_per_element,operations,Bands.plan(~Element,bands))