tensor_band_schedule.bend

stelliferous@0.0.2.0 · 78 lines · 15 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.2.0/tensor_band.bend as Bands
6import stelliferous@0.0.2.0/storage_buffer.bend as Storage
7import stelliferous@0.0.2.0/tensor_operations.bend as Operations
8import stelliferous@0.0.2.0/kernel_cost.bend as Cost
9import stelliferous@0.0.2.0/default_kernel_cost.bend as DefaultCost
10
11type Size is Data:
12  Size{leaves: U32,minimum_rows: U32,maximum_rows: U32}
13
14def merge(left: Size,right: Size) -> Size:
15  Size{ln,lmin,lmax} = left
16  Size{rn,rmin,rmax} = right
17  Size{(ln + rn : U32),U32.min(lmin,rmin),U32.max(lmax,rmax)}
18
19def size(layout: Bands.Plan) -> Size:
20  match layout:
21    case Bands.Rows{first,+rows}: Size{1,rows,rows}
22    case Bands.Split{left,right}: merge(size(left),size(right))
23
24# The whole operation, not each leaf, must reach the grain: small leaves only
25# lengthen the runtime's rounds, which the region estimate already counts.
26def sized(geometry: Bands.Geometry,+workers: U32,element_ns: F32,config: Cost.Configuration,size: Size) -> Bool:
27  Bands.Geometry{batch,channels,height,width} = geometry
28  Cost.Configuration{mac,byte_ns,operation_ns,grain,wake} = config
29  Size{+leaves,rows,maximum} = size
30  +row_ns = (U32.to_f32(channels) * U32.to_f32(width) * element_ns : F32)
31  +total = (U32.to_f32(height) * row_ns : F32)
32  U32.is_eq(batch,1) && U32.is_lt(1,workers) && U32.is_lt(1,leaves) && U32.is_le(leaves,(U32.min(workers,858993459) * 5 : U32)) &&
33    F32.is_le(grain,total) &&
34    F32.is_lt((Cost.region_ns(workers,leaves,(U32.to_f32(maximum) * row_ns : F32)) + 2.0 * wake : F32),total)
35
36def parallel(geometry: Bands.Geometry,workers: U32,bytes_per_element: F32,operations: F32,layout: Bands.Plan) -> Bool:
37  +config = DefaultCost.configuration()
38  sized(geometry,workers,Cost.element_ns(bytes_per_element,operations,config),config,size(layout))
39
40# For equal-shaped operands the logical read/write volume of repartition is
41# identical. The target owners' rounded capacities decide the remaining cost.
42def allocation_elements(layout: Bands.Plan,+geometry: Bands.Geometry) -> Nat:
43  match layout:
44    case Bands.Rows{first,rows}: U32.to_nat(U32.shln(1,Storage.depth_for(Bands.local_count(geometry,rows))))
45    case Bands.Split{left,right}: Nat.add(allocation_elements(left,geometry),allocation_elements(right,geometry))
46
47def alignment(+geometry: Bands.Geometry,+left: Bands.Plan,+right: Bands.Plan) -> Bands.Plan:
48  Bool.pick(Bands.Plan,Bands.same_plan(left,right) || Nat.is_le(allocation_elements(left,geometry),allocation_elements(right,geometry)),left,right)
49
50def 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):
51  (bands,+layout) = owners
52  Bands.map_with(~Element,~Context,~Bands.map_whole_context(~Element,~Context,~mapper),layout,parallel(geometry,workers,8.0,operations,layout),geometry,context,bands)
53
54def 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):
55  mapped_context(~Element,~Context,~mapper,context,geometry,workers,operations,Bands.plan(~Element,bands))
56
57def 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):
58  (bands,+layout) = owners
59  Bands.scalar(~Element,~operation,layout,parallel(geometry,workers,8.0,operations,layout),geometry,value,bands)
60
61def scalar(~Element: Data,~operation: Element -> Element -> Element,geometry: Bands.Geometry,workers: U32,operations: F32,value: Element,bands: Bands.Bands(~Element)) -> Bands.Bands(~Element):
62  scalar_planned(~Element,~operation,geometry,workers,operations,value,Bands.plan(~Element,bands))
63
64def 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):
65  (bands,+layout) = owners
66  Bands.apply_indexed(~Element,~operation,layout,parallel(geometry,workers,12.0,operations,layout),geometry,fill,bands)
67
68def indexed(~Element: Data,~operation: U32 -> Element -> Element,geometry: Bands.Geometry,workers: U32,operations: F32,fill: Element,bands: Bands.Bands(~Element)) -> Bands.Bands(~Element):
69  indexed_planned(~Element,~operation,geometry,workers,operations,fill,Bands.plan(~Element,bands))
70
71def transformed(~Element: Data,~Context: Data,~leaf: Context -> U32 -> U32 -> Storage.Buffer<Element> -> Bands.Bands(~Element),
72  context: Context,+geometry: Bands.Geometry,workers: U32,bytes_per_element: F32,owners: Bands.Bands(~Element) & Bands.Plan) -> Bands.Bands(~Element):
73  (bands,+layout) = owners
74  Bands.transform(~Element,~Context,~leaf,layout,parallel(geometry,workers,bytes_per_element,0.0,layout),geometry,context,bands)
75
76def transform(~Element: Data,~Context: Data,~leaf: Context -> U32 -> U32 -> Storage.Buffer<Element> -> Bands.Bands(~Element),
77  context: Context,geometry: Bands.Geometry,workers: U32,bytes_per_element: F32,bands: Bands.Bands(~Element)) -> Bands.Bands(~Element):
78  transformed(~Element,~Context,~leaf,context,geometry,workers,bytes_per_element,Bands.plan(~Element,bands))