1	# Concatenation: the parts meet in one plan of the rows of their joining and
2	# arrive in it. Every band of the output is the leaf concat of the bands of
3	# the parts there, along the middle axis of [outer, extent, inner]: one
4	# checked call per band.
5	import Base
6	import ./tensor.bend as Tensors
7	import ./tensor_band.bend as Bands
8	import ./storage_buffer.bend as Storage
9	import ./traversal_partition.bend as Partition
10	import ./tensor_band_schedule.bend as BandSchedule
11	
12	# How the parts lie in bands. Along the channels of images a band holds rows
13	# of one image. Along any other axis the tensor is [lines, extent, inner] and
14	# a band holds whole lines.
15	type Joining is Data:
16	  Channels{axis: U32,batch: U32,height: U32,width: U32}
17	  Lines{axis: U32,lines: U32,inner: U32}
18	
19	def lined(+shape: List<&2,U32>,+axis: U32) -> Joining:
20	  Lines{axis,Tensors.count_or_zero(List.take(&2,U32,shape,U32.to_nat(axis))),Tensors.count_or_zero(List.drop(&2,U32,shape,1n+U32.to_nat(axis)))}
21	
22	# The channels of an image shape, when the axis is its channel axis.
23	def imaged(image: Bool,+channels: Joining,+shape: List<&2,U32>,+axis: U32) -> Joining:
24	  match image:
25	    case True{}: channels
26	    case False{}: lined(shape,axis)
27	
28	def joining(shape: List<&2,U32>,+axis: U32) -> Joining:
29	  match shape:
30	    case [+channels,+height,+width]: imaged(U32.is_eq(axis,0),Channels{axis,1,height,width},[channels,height,width],axis)
31	    case [+batch,+channels,+height,+width]: imaged(U32.is_eq(axis,1),Channels{axis,batch,height,width},[batch,channels,height,width],axis)
32	    case +shape: lined(shape,axis)
33	
34	def axis(joining: Joining) -> U32:
35	  match joining:
36	    case Channels{axis,batch,height,width}: axis
37	    case Lines{axis,lines,inner}: axis
38	
39	# The extent of a part along the axis.
40	def extent(joining: Joining,shape: List<&2,U32>) -> U32:
41	  Tensors.known_count(List.get(&2,U32,shape,U32.to_nat(axis(joining))))
42	
43	# The geometry of a tensor of an extent along the axis.
44	def geometry(joining: Joining,+extent: U32) -> Bands.Geometry:
45	  match joining:
46	    case Channels{axis,batch,height,width}: Bands.Geometry{batch,extent,height,width}
47	    case Lines{axis,lines,inner}: Bands.Geometry{1,1,lines,(extent * inner : U32)}
48	
49	# A band of rows rows as [outer, extent, inner]: the rows of one image in
50	# every channel, or whole lines.
51	def band_outer(joining: Joining,+rows: U32) -> U32:
52	  match joining:
53	    case Channels{axis,batch,height,width}: 1
54	    case Lines{axis,lines,inner}: rows
55	
56	def band_inner(joining: Joining,+rows: U32) -> U32:
57	  match joining:
58	    case Channels{axis,batch,height,width}: (rows * width : U32)
59	    case Lines{axis,lines,inner}: inner
60	
61	def band_shape(+joining: Joining,+extent: U32,+rows: U32) -> List<&2,U32>:
62	  [band_outer(joining,rows),extent,band_inner(joining,rows)]
63	
64	# A part offers its plan when its bands are in the geometry of the joining,
65	# and one band of every image otherwise.
66	def part_target(rows: Bool,+geometry: Bands.Geometry,+layout: Bands.Plan) -> BandSchedule.Target:
67	  match rows:
68	    case True{}: BandSchedule.Target{geometry,layout,BandSchedule.splits(layout)}
69	    case False{}: BandSchedule.Target{geometry,Bands.unsplit(geometry),False{}}
70	
71	def part_offered(~Element: Data,+target: Bands.Geometry,+geometry: Bands.Geometry,+shape: List<&2,U32>,fill: Element,status: Tensors.Status,workers: U32,
72	  owners: Bands.Bands(~Element) & Bands.Plan) -> Tensors.Tensor<Element> & BandSchedule.Target:
73	  (bands,layout) = owners
74	  (Tensors.Tensor{bands,geometry,shape,fill,status,workers},part_target(Bands.same_geometry(geometry,target),target,layout))
75	
76	def part_offer(~Element: Data,+joining: Joining,tensor: Tensors.Tensor<Element>) -> Tensors.Tensor<Element> & BandSchedule.Target:
77	  Tensors.Tensor{bands,+current,+shape,fill,status,workers} = tensor
78	  part_offered(~Element,geometry(joining,extent(joining,shape)),current,shape,fill,status,workers,Bands.plan(~Element,bands))
79	
80	def offered_single(~Element: Data,offered: Tensors.Tensor<Element> & BandSchedule.Target) -> Tensors.Parts<Element> & BandSchedule.Target:
81	  (tensor,target) = offered
82	  (Tensors.Single{tensor},target)
83	
84	def offered_cons(~Element: Data,offered: Tensors.Tensor<Element> & BandSchedule.Target,rest: Tensors.Parts<Element> & BandSchedule.Target) ->
85	  Tensors.Parts<Element> & BandSchedule.Target:
86	  (tensor,target) = offered
87	  (others,other_target) = rest
88	  (Tensors.Cons{tensor,others},BandSchedule.meet(target,other_target))
89	
90	# The parts in order beside the layout they meet in.
91	def offers(~Element: Data,+joining: Joining,parts: Tensors.Parts<Element>) -> Tensors.Parts<Element> & BandSchedule.Target:
92	  match parts:
93	    case Tensors.Single{tensor}: offered_single(~Element,part_offer(~Element,joining,tensor))
94	    case Tensors.Cons{tensor,rest}: offered_cons(~Element,part_offer(~Element,joining,tensor),offers(~Element,joining,rest))
95	
96	# A part in the geometry of the joining, cut by the plan the parts met in.
97	def arrived(~Element: Data,+joining: Joining,+layout: Bands.Plan,tensor: Tensors.Tensor<Element>) -> Tensors.Tensor<Element>:
98	  Tensors.Tensor{bands,+current,+shape,+fill,+status,+workers} = tensor
99	  +target = geometry(joining,extent(joining,shape))
100	  Tensors.Tensor{Tensors.arrive(~Element,target,layout,Tensors.Tensor{bands,current,shape,fill,status,workers}),target,shape,fill,status,workers}
101	
102	def arrived_parts(~Element: Data,+joining: Joining,+layout: Bands.Plan,parts: Tensors.Parts<Element>) -> Tensors.Parts<Element>:
103	  match parts:
104	    case Tensors.Single{tensor}: Tensors.Single{arrived(~Element,joining,layout,tensor)}
105	    case Tensors.Cons{tensor,rest}: Tensors.Cons{arrived(~Element,joining,layout,tensor),arrived_parts(~Element,joining,layout,rest)}
106	
107	# ---- The bands of the parts, band by band ----
108	
109	# The bands of every part at one band of the output, as whole tensors of their
110	# band shapes in the order of the parts, and whether every part has a band of
111	# the same rows there.
112	type Gathered<-Element: Data> is Type:
113	  Gathered{parts: Tensors.Parts<Element>,fill: Element,valid: Bool}
114	
115	def Tree(~Element: Data) -> Type:
116	  Partition.Tree<&1,Gathered<Element>,Unit>
117	
118	def band_tensor(~Element: Data,+joining: Joining,+extent: U32,+fill: Element,+rows: U32,buffer: Storage.Buffer<Element>) -> Tensors.Tensor<Element>:
119	  Tensors.whole(~Element,buffer,band_shape(joining,extent,rows),fill,Tensors.Usable{},1)
120	
121	def last_bands(~Element: Data,+joining: Joining,+extent: U32,+fill: Element,bands: Bands.Bands(~Element)) -> Tree(~Element):
122	  match bands:
123	    case Partition.Leaf{buffer,+first,+rows,details}:
124	      Partition.Leaf{Gathered{Tensors.Single{band_tensor(~Element,joining,extent,fill,rows,buffer)},fill,True{}},first,rows,Unit{}}
125	    case Partition.Fork{left,right}:
126	      Partition.Fork{last_bands(~Element,joining,extent,fill,left),last_bands(~Element,joining,extent,fill,right)}
127	
128	def invalid_band(~Element: Data,gathered: Gathered<Element>,first: U32,rows: U32) -> Tree(~Element):
129	  Gathered{parts,fill,valid} = gathered
130	  Partition.Leaf{Gathered{parts,fill,False{}},first,rows,Unit{}}
131	
132	def invalid(~Element: Data,gathered: Tree(~Element)) -> Tree(~Element):
133	  match gathered:
134	    case Partition.Leaf{leaf,first,rows,details}: invalid_band(~Element,leaf,first,rows)
135	    case Partition.Fork{left,right}: Partition.Fork{invalid(~Element,left),invalid(~Element,right)}
136	
137	def joined_band(~Element: Data,+joining: Joining,+extent: U32,+fill: Element,buffer: Storage.Buffer<Element>,+first: U32,+rows: U32,
138	  gathered: Gathered<Element>,+other_first: U32,+other_rows: U32) -> Tree(~Element):
139	  Gathered{parts,other_fill,valid} = gathered
140	  Partition.Leaf{Gathered{Tensors.Cons{band_tensor(~Element,joining,extent,fill,rows,buffer),parts},other_fill,
141	    valid && U32.is_eq(first,other_first) && U32.is_eq(rows,other_rows)},first,rows,Unit{}}
142	
143	# The bands of one more part in front: two trees of different forms, or two
144	# bands of different rows, leave the band invalid.
145	def joined_bands(~Element: Data,+joining: Joining,+extent: U32,+fill: Element,bands: Bands.Bands(~Element),gathered: Tree(~Element)) -> Tree(~Element):
146	  match bands gathered:
147	    case Partition.Leaf{buffer,first,rows,details} Partition.Leaf{other,other_first,other_rows,other_details}:
148	      joined_band(~Element,joining,extent,fill,buffer,first,rows,other,other_first,other_rows)
149	    case Partition.Fork{left,right} Partition.Fork{other_left,other_right}:
150	      Partition.Fork{joined_bands(~Element,joining,extent,fill,left,other_left),joined_bands(~Element,joining,extent,fill,right,other_right)}
151	    case bands gathered: invalid(~Element,gathered)
152	
153	def last_part(~Element: Data,+joining: Joining,tensor: Tensors.Tensor<Element>) -> Tree(~Element):
154	  Tensors.Tensor{bands,geometry,+shape,+fill,status,workers} = tensor
155	  last_bands(~Element,joining,extent(joining,shape),fill,bands)
156	
157	def joined_part(~Element: Data,+joining: Joining,tensor: Tensors.Tensor<Element>,gathered: Tree(~Element)) -> Tree(~Element):
158	  Tensors.Tensor{bands,geometry,+shape,+fill,status,workers} = tensor
159	  joined_bands(~Element,joining,extent(joining,shape),fill,bands,gathered)
160	
161	# One tree whose bands hold the bands of every part, without copying values.
162	def gather(~Element: Data,+joining: Joining,parts: Tensors.Parts<Element>) -> Tree(~Element):
163	  match parts:
164	    case Tensors.Single{tensor}: last_part(~Element,joining,tensor)
165	    case Tensors.Cons{tensor,rest}: joined_part(~Element,joining,tensor,gather(~Element,joining,rest))
166	
167	# ---- One checked leaf concat per band ----
168	
169	type Context is Data:
170	  Context{joining: Joining,extent: U32}
171	
172	def checked(~Element: Data,valid: Bool,+first: U32,+rows: U32,tensor: Tensors.Tensor<Element>) -> Bands.Transformed(~Element):
173	  Tensors.with_whole(~Element,~Bands.Transformed(~Element),tensor,storage => shape => fill => status => workers =>
174	    Partition.Leaf{Bands.Checked{storage,valid && Tensors.usable(status)},first,rows,Unit{}})
175	
176	def band(~Element: Data,~joined: Tensors.Parts<Element> -> List<&2,U32> -> Element -> Tensors.Tensor<Element>,
177	  context: Context,first: U32,rows: U32,gathered: Gathered<Element>,resource: Unit) -> Bands.Transformed(~Element):
178	  Context{joining,extent} = context
179	  Gathered{parts,fill,valid} = gathered
180	  +rows = rows
181	  checked(~Element,valid,first,rows,joined(parts,band_shape(joining,extent,rows),fill))
182	
183	def unjoined_band(~Element: Data,gathered: Gathered<Element>,first: U32,rows: U32) -> Bands.Transformed(~Element):
184	  Gathered{parts,+fill,valid} = gathered
185	  Partition.Leaf{Bands.Checked{Tensors.buffer(~Element,fill,0),False{}},first,rows,Unit{}}
186	
187	def unjoined(~Element: Data,gathered: Tree(~Element)) -> Bands.Transformed(~Element):
188	  match gathered:
189	    case Partition.Leaf{leaf,first,rows,details}: unjoined_band(~Element,leaf,first,rows)
190	    case Partition.Fork{left,right}: Partition.Fork{unjoined(~Element,left),unjoined(~Element,right)}
191	
192	# Every band of the output through the leaf concat joined of the bands of
193	# the parts, beside whether every call was usable.
194	def execute(~Element: Data,~joined: Tensors.Parts<Element> -> List<&2,U32> -> Element -> Tensors.Tensor<Element>,
195	  +layout: Bands.Plan,+joining: Joining,+extent: U32,+workers: U32,parts: Tensors.Parts<Element>) -> Bands.Bands(~Element) & Bool:
196	  +output = geometry(joining,extent)
197	  +parallel = BandSchedule.parallel(output,workers,8.0,0.0,layout)
198	  Bands.reduce(~Bands.Checked<Element>,~Storage.Buffer<Element>,~Bool,~Unit,~Bands.checked_band(~Element),~Bool.and,Unit{},
199	    Bands.each(~Gathered<Element>,~Bands.Checked<Element>,~Unit,~Context,~Bands.unshared,~band(~Element,~joined),~unjoined(~Element),
200	      layout,parallel,Context{joining,extent},gather(~Element,joining,parts),Unit{}))
201	
202	# The bands of the output when every band call was usable and the bands
203	# cover the output; otherwise concat fails with the output shape.
204	def join(~Element: Data,~joined: Tensors.Parts<Element> -> List<&2,U32> -> Element -> Tensors.Tensor<Element>,
205	  +layout: Bands.Plan,+joining: Joining,+extent: U32,parts: Tensors.Parts<Element>,+shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
206	  Tensors.transformed(~Element,"concat",Tensors.InvalidArgument{shape},geometry(joining,extent),shape,fill,workers,execute(~Element,~joined,layout,joining,extent,workers,parts))
207	
208	# Every part arrives in the plan the parts met in, in the geometry of the
209	# joining, and the band concat joins them.
210	def met(~Element: Data,~joined: Tensors.Parts<Element> -> List<&2,U32> -> Element -> Tensors.Tensor<Element>,
211	  +joining: Joining,+shape: List<&2,U32>,+fill: Element,+workers: U32,offered: Tensors.Parts<Element> & BandSchedule.Target) -> Tensors.Tensor<Element>:
212	  (parts,target) = offered
213	  BandSchedule.Target{other,+layout,split} = target
214	  join(~Element,~joined,layout,joining,extent(joining,shape),arrived_parts(~Element,joining,layout,parts),shape,fill,workers)