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.
4	import Base
5	import ./tensor_band.bend as Bands
6	import ./storage_buffer.bend as Storage
7	import ./kernel_cost.bend as Cost
8	import ./default_kernel_cost.bend as DefaultCost
9	
10	type Size is Data:
11	  Size{leaves: U32,minimum_rows: U32,maximum_rows: U32}
12	
13	def 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	
18	def leaf_size(context: Unit,first: U32,rows: U32) -> Size:
19	  +rows = rows
20	  Size{1,rows,rows}
21	
22	def 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.
27	def 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	
37	def 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.
43	def 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	
46	def 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	
49	def 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.
54	type Target is Data:
55	  Target{geometry: Bands.Geometry,layout: Bands.Plan,split: Bool}
56	
57	def splits(layout: Bands.Plan) -> Bool:
58	  match layout:
59	    case Bands.Rows{first,count}: False{}
60	    case Bands.Split{left,right}: True{}
61	
62	def 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	
67	def 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.
77	def 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	
82	def 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.
90	def 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	
94	def 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	
98	def 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	
101	def 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	
105	def 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	
108	def 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	
112	def 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	
115	def 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.
122	def 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))