1	# Window geometry is independent of convolution arithmetic. Each output row
2	# belongs to one existing core band; neighbours contribute only halo rows.
3	import Base
4	import ./tensor_band.bend as Bands
5	import ./storage_buffer.bend as Storage
6	import ./traversal_partition.bend as Partition
7	
8	type Geometry is Data:
9	  Geometry{input: Bands.Geometry,kernel: U32,stride: U32,padding: U32,output_height: U32,output_width: U32}
10	
11	type Rows is Data:
12	  Rows{first: U32,count: U32}
13	
14	# A leaf's output rows, its halo rows above and below the core, and the
15	# image rows its window holds: from the stride-aligned first row its outputs
16	# read to the last, so the window packs as a shorter image.
17	type RowWindow is Data:
18	  RowWindow{output: Rows,above: Rows,below: Rows,window: Rows}
19	
20	def ceiling(+value: U32,+stride: U32) -> U32:
21	  (value / stride + Bool.pick(U32,U32.is_zero((value % stride : U32)),0,1) : U32)
22	
23	def subtract_clipped(+value: U32,+amount: U32) -> U32:
24	  (U32.max(value,amount) - amount : U32)
25	
26	def clipped(+first: U32,+end: U32) -> Rows:
27	  Rows{first,subtract_clipped(end,first)}
28	
29	def nonempty(active: Bool,window: RowWindow) -> RowWindow:
30	  match active window:
31	    case True{} window: window
32	    case False{} RowWindow{output,Rows{above_first,above_count},Rows{below_first,below_count},Rows{window_first,window_count}}:
33	      RowWindow{output,Rows{above_first,0},Rows{below_first,0},Rows{window_first,0}}
34	
35	# The first row a window holds: the first row its outputs read, moved down
36	# to a stride multiple.
37	def start(+height: U32,+output_first: U32,+stride: U32,+padding: U32) -> U32:
38	  +needed = U32.min(height,subtract_clipped((output_first * stride : U32),padding))
39	  subtract_clipped(needed,(needed % stride : U32))
40	
41	# Output partition boundaries follow the input cores, with the first and last
42	# cores also owning any output rows introduced by padding. The halos hold
43	# every window row outside the core, so a window is its image rows exactly.
44	# Caller geometry checks establish positive stride/height and bounded products.
45	def window(geometry: Geometry,+first: U32,+count: U32) -> RowWindow:
46	  Geometry{Bands.Geometry{batch,channels,+height,width},+kernel,+stride,+padding,+output_height,output_width} = geometry
47	  +image = (first / height : U32)
48	  +core_first = (first % height : U32)
49	  +core_end = (core_first + count : U32)
50	  +out_first = U32.min(output_height,ceiling(core_first,stride))
51	  +out_end = Bool.pick(U32,U32.is_eq(core_end,height),output_height,U32.min(output_height,ceiling(core_end,stride)))
52	  +needed_first = start(height,out_first,stride,padding)
53	  +needed_end = U32.min(height,subtract_clipped((subtract_clipped(out_end,1) * stride + kernel : U32),padding))
54	  +above_first = (image * height + needed_first : U32)
55	  +above_end = (image * height + U32.min(core_first,needed_end) : U32)
56	  +below_first = (image * height + U32.max(core_end,needed_first) : U32)
57	  +below_end = (image * height + needed_end : U32)
58	  nonempty(U32.is_lt(out_first,out_end),RowWindow{Rows{(image * output_height + out_first : U32),subtract_clipped(out_end,out_first)},
59	    clipped(above_first,above_end),clipped(below_first,below_end),clipped(above_first,below_end)})
60	
61	type Halo<-Element: Data> is Type:
62	  Halo{above: Storage.Buffer<Element>,below: Storage.Buffer<Element>,window: RowWindow}
63	
64	def Halos(~Element: Data) -> Type:
65	  Partition.Tree<&1,Halo<Element>,Unit>
66	
67	def extract(~Element: Data,+geometry: Bands.Geometry,fill: Element,range: Rows,bands: Bands.Bands(~Element)) -> Bands.Bands(~Element) & Storage.Buffer<Element>:
68	  Rows{+first,+count} = range
69	  Bands.range(~Element,geometry,first,count,bands,Bands.allocate(~Element,fill,Bands.local_count(geometry,count)))
70	
71	def halo(~Element: Data,+geometry: Bands.Geometry,+fill: Element,+first: U32,+count: U32,window: RowWindow,bands: Bands.Bands(~Element)) -> Bands.Bands(~Element) & Halos(~Element):
72	  RowWindow{+output,+above,+below,+rows} = window
73	  Storage.bind_pair(~Bands.Bands(~Element),~Storage.Buffer<Element>,~(Bands.Bands(~Element) & Halos(~Element)),
74	    extract(~Element,geometry,fill,above,bands),bands => upper =>
75	      Storage.bind_pair(~Bands.Bands(~Element),~Storage.Buffer<Element>,~(Bands.Bands(~Element) & Halos(~Element)),
76	        extract(~Element,geometry,fill,below,bands),bands => lower =>
77	          (bands,Partition.Leaf{Halo{upper,lower,RowWindow{output,above,below,rows}},first,count,Unit{}})))
78	
79	def copied_window(~Element: Data,owners: Bands.Bands(~Element) & Storage.Buffer<Element>) -> Storage.Buffer<Element>:
80	  (sources,window) = owners
81	  window
82	
83	# The window rows of every channel, stored contiguously. A window equal to
84	# the core (for example a 1x1 stride-1 operation) is the core itself; any
85	# other window is the shared range copy of its rows from the upper halo, the
86	# core and the lower halo, one contiguous range per channel per source. The
87	# three copies are direct calls: a leaf task must not reach a closure.
88	def assembled(~Element: Data,exact: Bool,+fill: Element,+geometry: Bands.Geometry,+window: Rows,+core: Rows,+above: Rows,+below: Rows,
89	  core_buffer: Storage.Buffer<Element>,above_buffer: Storage.Buffer<Element>,below_buffer: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
90	  match exact:
91	    case True{}: core_buffer
92	    case False{}:
93	      Rows{+window_first,+window_rows} = window
94	      Rows{core_first,core_rows} = core
95	      Rows{above_first,above_rows} = above
96	      Rows{below_first,below_rows} = below
97	      copied_window(~Element,Bands.range_leaf(~Element,geometry,window_first,window_rows,0,below_first,below_rows,below_buffer,
98	        copied_window(~Element,Bands.range_leaf(~Element,geometry,window_first,window_rows,0,core_first,core_rows,core_buffer,
99	          copied_window(~Element,Bands.range_leaf(~Element,geometry,window_first,window_rows,0,above_first,above_rows,above_buffer,
100	            Bands.allocate(~Element,fill,Bands.local_count(geometry,window_rows))))))))
101	
102	def assemble(~Element: Data,+fill: Element,+geometry: Bands.Geometry,+window: Rows,+core: Rows,+above: Rows,+below: Rows,
103	  core_buffer: Storage.Buffer<Element>,above_buffer: Storage.Buffer<Element>,below_buffer: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
104	  Rows{+window_first,+window_rows} = window
105	  Rows{+core_first,+core_rows} = core
106	  assembled(~Element,U32.is_eq(window_first,core_first) && U32.is_eq(window_rows,core_rows),fill,geometry,window,core,above,below,
107	    core_buffer,above_buffer,below_buffer)
108	
109	def input_geometry(geometry: Geometry) -> Bands.Geometry:
110	  Geometry{input,kernel,stride,padding,output_height,output_width} = geometry
111	  input
112	
113	# The full core tree is threaded through the preparation pass without cloning.
114	# Each allocation is a requested neighbour range, never a core-sized window.
115	# The resulting halo tree follows the same layout, ready to pair with cores.
116	def prepare(~Element: Data,+geometry: Geometry,+fill: Element,layout: Bands.Plan,bands: Bands.Bands(~Element)) -> Bands.Bands(~Element) & Halos(~Element):
117	  match layout:
118	    case Bands.Rows{+first,+count}: halo(~Element,input_geometry(geometry),fill,first,count,window(geometry,first,count),bands)
119	    case Bands.Split{left,right}:
120	      Storage.bind_pair(~Bands.Bands(~Element),~Halos(~Element),~(Bands.Bands(~Element) & Halos(~Element)),
121	        prepare(~Element,geometry,fill,left,bands),bands => left =>
122	          Storage.bind_pair(~Bands.Bands(~Element),~Halos(~Element),~(Bands.Bands(~Element) & Halos(~Element)),
123	            prepare(~Element,geometry,fill,right,bands),bands => right => (bands,Partition.Fork{left,right})))
124	
125	def window_band(~Element: Data,~Context: Data,~leaf: Context -> U32 -> U32 -> Storage.Buffer<Element> -> Halo<Element> -> Bands.Bands(~Element),
126	  context: Context,first: U32,rows: U32,input: Storage.Buffer<Element>,halo: Halo<Element>,resource: Unit) -> Bands.Bands(~Element):
127	  leaf(context,first,rows,input,halo)
128	
129	# Every core leaf with its halo, through the shared band recursion. Trees of
130	# another shape keep the cores.
131	def apply(~Element: Data,~Context: Data,~leaf: Context -> U32 -> U32 -> Storage.Buffer<Element> -> Halo<Element> -> Bands.Bands(~Element),
132	  layout: Bands.Plan,mode: Bool,enabled: Bool,geometry: Bands.Geometry,context: Context,bands: Bands.Bands(~Element),halos: Halos(~Element)) -> Bands.Bands(~Element):
133	  Bands.each(~Bands.Paired<Storage.Buffer<Element>,Halo<Element>>,~Storage.Buffer<Element>,~Unit,~Context,~Bands.unshared,
134	    ~Bands.paired(~Storage.Buffer<Element>,~Halo<Element>,~Unit,~Context,~window_band(~Element,~Context,~leaf)),
135	    ~Bands.left_of(~Storage.Buffer<Element>,~Halo<Element>),
136	    layout,mode,enabled,geometry,context,Bands.zip(~Storage.Buffer<Element>,~Halo<Element>,bands,halos),Unit{})