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	# What the halo of one band needs.
117	type Padded<-Element: Data> is Data:
118	  Padded{geometry: Geometry,fill: Element}
119	
120	def halo_band(~Element: Data,context: Padded<Element>,first: U32,count: U32,bands: Bands.Bands(~Element)) -> Bands.Bands(~Element) & Halos(~Element):
121	  Padded{+geometry,fill} = context
122	  +first = first
123	  +count = count
124	  halo(~Element,input_geometry(geometry),fill,first,count,window(geometry,first,count),bands)
125	
126	def prepare(~Element: Data,+geometry: Geometry,+fill: Element,layout: Bands.Plan,bands: Bands.Bands(~Element)) -> Bands.Bands(~Element) & Halos(~Element):
127	  Bands.build(~Halo<Element>,~Bands.Bands(~Element),~Padded<Element>,~halo_band(~Element),layout,Padded{geometry,fill},bands)
128	
129	# A core with its halo goes through the leaf; a core without one is not
130	# transformed.
131	def window_band(~Element: Data,~Context: Data,~leaf: Context -> U32 -> U32 -> Storage.Buffer<Element> -> Halo<Element> -> Bands.Transformed(~Element),
132	  context: Context,first: U32,rows: U32,pair: Bands.Paired<Storage.Buffer<Element>,Halo<Element>>,resource: Unit) -> Bands.Transformed(~Element):
133	  match pair:
134	    case Bands.Both{input,halo}: leaf(context,first,rows,input,halo)
135	    case Bands.Only{input,own_first,own_rows}: Partition.Leaf{Bands.Checked{input,False{}},own_first,own_rows,Unit{}}
136	
137	def unwindowed(~Element: Data,tree: Partition.Tree<&1,Bands.Paired<Storage.Buffer<Element>,Halo<Element>>,Unit>) -> Bands.Transformed(~Element):
138	  Bands.untransformed(~Element,Bands.left_of(~Storage.Buffer<Element>,~Halo<Element>,tree))
139	
140	# Every core leaf with its halo through its own checked call, in the shared
141	# band recursion. The bands come back beside whether every call was usable;
142	# trees of another shape are not transformed.
143	def apply(~Element: Data,~Context: Data,~leaf: Context -> U32 -> U32 -> Storage.Buffer<Element> -> Halo<Element> -> Bands.Transformed(~Element),
144	  layout: Bands.Plan,enabled: Bool,context: Context,bands: Bands.Bands(~Element),halos: Halos(~Element)) ->
145	  Bands.Bands(~Element) & Bool:
146	  Bands.reduce(~Bands.Checked<Element>,~Storage.Buffer<Element>,~Bool,~Unit,~Bands.checked_band(~Element),~Bool.and,Unit{},
147	    Bands.each(~Bands.Paired<Storage.Buffer<Element>,Halo<Element>>,~Bands.Checked<Element>,~Unit,~Context,~Bands.unshared,
148	      ~window_band(~Element,~Context,~leaf),~unwindowed(~Element),
149	      layout,enabled,context,Bands.zip(~Storage.Buffer<Element>,~Halo<Element>,bands,halos),Unit{}))