tensor_band_concat.bend
stelliferous@0.0.3.0 · 214 lines · 0 definitions
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.
5import Base
6import stelliferous@0.0.3.0/tensor.bend as Tensors
7import stelliferous@0.0.3.0/tensor_band.bend as Bands
8import stelliferous@0.0.3.0/storage_buffer.bend as Storage
9import stelliferous@0.0.3.0/traversal_partition.bend as Partition
10import stelliferous@0.0.3.0/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.
15type Joining is Data:
16 Channels{axis: U32,batch: U32,height: U32,width: U32}
17 Lines{axis: U32,lines: U32,inner: U32}
18
19def 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.
23def 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
28def 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
34def 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.
40def 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.
44def 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.
51def 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
56def 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
61def 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.
66def 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
71def 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
76def 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
80def 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
84def 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.
91def 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.
97def 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
102def 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.
112type Gathered<-Element: Data> is Type:
113 Gathered{parts: Tensors.Parts<Element>,fill: Element,valid: Bool}
114
115def Tree(~Element: Data) -> Type:
116 Partition.Tree<&1,Gathered<Element>,Unit>
117
118def 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
121def 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
128def 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
132def 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
137def 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.
145def 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
153def 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
157def 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.
162def 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
169type Context is Data:
170 Context{joining: Joining,extent: U32}
171
172def 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
176def 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
183def 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
187def 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.
194def 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.
204def 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.
210def 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)