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 ./tensor_operations.bend as Operations
8	import ./kernel_cost.bend as Cost
9	import ./default_kernel_cost.bend as DefaultCost
10	
11	type Size is Data:
12	  Size{leaves: U32,minimum_rows: U32,maximum_rows: U32}
13	
14	def 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	
19	def 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.
26	def 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	
36	def 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.
42	def 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	
47	def 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	
50	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):
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	
54	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):
55	  mapped_context(~Element,~Context,~mapper,context,geometry,workers,operations,Bands.plan(~Element,bands))
56	
57	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):
58	  (bands,+layout) = owners
59	  Bands.scalar(~Element,~operation,layout,parallel(geometry,workers,8.0,operations,layout),geometry,value,bands)
60	
61	def 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	
64	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):
65	  (bands,+layout) = owners
66	  Bands.apply_indexed(~Element,~operation,layout,parallel(geometry,workers,12.0,operations,layout),geometry,fill,bands)
67	
68	def 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	
71	def 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	
76	def 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))