tensor_band_concat.bend

stelliferous@0.0.2.0 · 181 lines · 31 definitions

1# Channel concatenation writes each input range directly into its final band.
2# Different input boundaries need no temporary dense tensor or aligned copy.
3import Base
4import stelliferous@0.0.2.0/tensor.bend as Tensors
5import stelliferous@0.0.2.0/tensor_band.bend as Bands
6import stelliferous@0.0.2.0/storage_buffer.bend as Storage
7import stelliferous@0.0.2.0/traversal_partition.bend as Partition
8import stelliferous@0.0.2.0/storage_copy.bend as Copy
9import stelliferous@0.0.2.0/tensor_band_schedule.bend as BandSchedule
10
11def geometry(shape: List<&2,U32>,axis: U32) -> Maybe<&2,Bands.Geometry>:
12  match shape axis:
13    case [channels,height,width] 0: Some{Bands.Geometry{1,channels,height,width}}
14    case [batch,channels,height,width] 1: Some{Bands.Geometry{batch,channels,height,width}}
15    case shape axis: None{}
16
17def tensor_planned(~Element: Data,geometry: Bands.Geometry,shape: List<&2,U32>,fill: Element,status: Tensors.Status,workers: U32,
18  owners: Bands.Bands(~Element) & Bands.Plan) -> Tensors.Tensor<Element> & Maybe<&2,Bands.Plan>:
19  (bands,layout) = owners
20  (Tensors.Banded{bands,geometry,shape,fill,status,workers},Some{layout})
21
22def tensor_plan(~Element: Data,tensor: Tensors.Tensor<Element>) -> Tensors.Tensor<Element> & Maybe<&2,Bands.Plan>:
23  match tensor:
24    case Tensors.Tensor{storage,shape,fill,status,workers}: (Tensors.Tensor{storage,shape,fill,status,workers},None{})
25    case Tensors.Banded{bands,geometry,shape,fill,status,workers}: tensor_planned(~Element,geometry,shape,fill,status,workers,Bands.plan(~Element,bands))
26
27def first_plan(first: Maybe<&2,Bands.Plan>,rest: Maybe<&2,Bands.Plan>) -> Maybe<&2,Bands.Plan>:
28  match first:
29    case Some{layout}: Some{layout}
30    case None{}: rest
31
32def single_plan(~Element: Data,first: Tensors.Tensor<Element> & Maybe<&2,Bands.Plan>) -> Tensors.Parts<Element> & Maybe<&2,Bands.Plan>:
33  (tensor,layout) = first
34  (Tensors.Single{tensor},layout)
35
36def cons_plan(~Element: Data,first: Tensors.Tensor<Element> & Maybe<&2,Bands.Plan>,rest: Tensors.Parts<Element> & Maybe<&2,Bands.Plan>) -> Tensors.Parts<Element> & Maybe<&2,Bands.Plan>:
37  (tensor,first) = first
38  (rest,other) = rest
39  (Tensors.Cons{tensor,rest},first_plan(first,other))
40
41def discover(~Element: Data,parts: Tensors.Parts<Element>) -> Tensors.Parts<Element> & Maybe<&2,Bands.Plan>:
42  match parts:
43    case Tensors.Single{tensor}: single_plan(~Element,tensor_plan(~Element,tensor))
44    case Tensors.Cons{tensor,rest}: cons_plan(~Element,tensor_plan(~Element,tensor),discover(~Element,rest))
45
46# Reuse the existing metadata traversal to classify the common band path.
47# Mixed representations and different plans use the dense concat caller.
48def compatible_plans(first: Maybe<&2,Bands.Plan>,rest: Maybe<&2,Bands.Plan>) -> Bool:
49  match first rest:
50    case Some{first} Some{rest}: Bands.same_plan(first,rest)
51    case first rest: False{}
52
53def single_compatible_plan(~Element: Data,tensor: Tensors.Tensor<Element>,layout: Maybe<&2,Bands.Plan>) ->
54  (Tensors.Parts<Element> & Maybe<&2,Bands.Plan>) & Bool:
55  match layout:
56    case Some{layout}: ((Tensors.Single{tensor},Some{layout}),True{})
57    case None{}: ((Tensors.Single{tensor},None{}),False{})
58
59def single_compatible(~Element: Data,first: Tensors.Tensor<Element> & Maybe<&2,Bands.Plan>) ->
60  (Tensors.Parts<Element> & Maybe<&2,Bands.Plan>) & Bool:
61  (tensor,layout) = first
62  single_compatible_plan(~Element,tensor,layout)
63
64def cons_compatible(~Element: Data,first: Tensors.Tensor<Element> & Maybe<&2,Bands.Plan>,
65  rest: (Tensors.Parts<Element> & Maybe<&2,Bands.Plan>) & Bool) -> (Tensors.Parts<Element> & Maybe<&2,Bands.Plan>) & Bool:
66  (tensor,first_layout) = first
67  (rest_and_layout,compatible) = rest
68  (rest_parts,rest_layout) = rest_and_layout
69  +first_layout = {first_layout : Maybe<&2,Bands.Plan>}
70  +rest_layout = {rest_layout : Maybe<&2,Bands.Plan>}
71  ((Tensors.Cons{tensor,rest_parts},first_plan(first_layout,rest_layout)),compatible && compatible_plans(first_layout,rest_layout))
72
73def discover_compatible(~Element: Data,parts: Tensors.Parts<Element>) -> (Tensors.Parts<Element> & Maybe<&2,Bands.Plan>) & Bool:
74  match parts:
75    case Tensors.Single{tensor}: single_compatible(~Element,tensor_plan(~Element,tensor))
76    case Tensors.Cons{tensor,rest}: cons_compatible(~Element,tensor_plan(~Element,tensor),discover_compatible(~Element,rest))
77
78def allocate(~Element: Data,layout: Bands.Plan,+geometry: Bands.Geometry,+fill: Element) -> Bands.Bands(~Element):
79  match layout:
80    case Bands.Rows{first,+rows}: Partition.Leaf{Tensors.buffer(~Element,fill,Bands.local_count(geometry,rows)),first,rows,Unit{}}
81    case Bands.Split{left,right}: Partition.Fork{allocate(~Element,left,geometry,fill),allocate(~Element,right,geometry,fill)}
82
83def dense_rows(~Element: Data,geometry: Bands.Geometry,+first: U32,+rows: U32,channel: U32,input: Storage.Buffer<Element>,output: Storage.Buffer<Element>) -> Storage.Buffer<Element> & Storage.Buffer<Element>:
84  Bands.Geometry{batch,+channels,+height,+width} = geometry
85  Bands.planes(~Element,U32.to_nat(channels),0,(rows * width : U32),(height * width : U32),(rows * width : U32),
86    ((first / height * channels * height + first % height) * width : U32),(channel * rows * width : U32),(input,output))
87
88def channel_offset(geometry: Bands.Geometry,channel: U32,rows: U32) -> U32:
89  Bands.Geometry{batch,channels,height,width} = geometry
90  (channel * rows * width : U32)
91
92def rows(~Element: Data,geometry: Bands.Geometry,+first: U32,+rows: U32,channel: U32,tensor: Tensors.Tensor<Element>,output: Storage.Buffer<Element>) -> Tensors.Tensor<Element> & Storage.Buffer<Element>:
93  match tensor:
94    case Tensors.Tensor{input,shape,fill,status,workers}:
95      Storage.bind_pair(~Storage.Buffer<Element>,~Storage.Buffer<Element>,~(Tensors.Tensor<Element> & Storage.Buffer<Element>),
96        dense_rows(~Element,geometry,first,rows,channel,input,output),input => output => (Tensors.Tensor{input,shape,fill,status,workers},output))
97    case Tensors.Banded{bands,+geometry,shape,fill,status,workers}:
98      Storage.bind_pair(~Bands.Bands(~Element),~Storage.Buffer<Element>,~(Tensors.Tensor<Element> & Storage.Buffer<Element>),
99        Bands.range_into(~Element,geometry,first,rows,channel_offset(geometry,channel,rows),bands,output),bands => output =>
100          (Tensors.Banded{bands,geometry,shape,fill,status,workers},output))
101
102def append_tree(~Element: Data,output: Bands.Bands(~Element),+geometry: Bands.Geometry,+channel: U32,input: Tensors.Tensor<Element>) -> Tensors.Tensor<Element> & Bands.Bands(~Element):
103  match output:
104    case Partition.Leaf{output,+first,+count,details}:
105      Storage.bind_pair(~Tensors.Tensor<Element>,~Storage.Buffer<Element>,~(Tensors.Tensor<Element> & Bands.Bands(~Element)),
106        rows(~Element,geometry,first,count,channel,input,output),input => output => (input,Partition.Leaf{output,first,count,Unit{}}))
107    case Partition.Fork{left,right}:
108      Storage.bind_pair(~Tensors.Tensor<Element>,~Bands.Bands(~Element),~(Tensors.Tensor<Element> & Bands.Bands(~Element)),
109        append_tree(~Element,left,geometry,channel,input),input => left =>
110          Storage.bind_pair(~Tensors.Tensor<Element>,~Bands.Bands(~Element),~(Tensors.Tensor<Element> & Bands.Bands(~Element)),
111            append_tree(~Element,right,geometry,channel,input),input => right => (input,Partition.Fork{left,right})))
112
113def channels(shape: List<&2,U32>) -> U32:
114  match shape:
115    case [channels,height,width]: channels
116    case [batch,channels,height,width]: channels
117    case shape: 0
118
119type Channel is Data:
120  Channel{geometry: Bands.Geometry,channel: U32}
121
122def append_band(~Element: Data,context: Channel,first: U32,rows: U32,output: Storage.Buffer<Element>,input: Storage.Buffer<Element>,resource: Unit) -> Bands.Bands(~Element):
123  Channel{+geometry,channel} = context
124  +rows = rows
125  Partition.Leaf{Bands.destination(~Element,Copy.range(~Element,Bands.local_count(geometry,rows),0,channel_offset(geometry,channel,rows),input,output)),first,rows,Unit{}}
126
127def append_planned(~Element: Data,+parallel: Bool,+geometry: Bands.Geometry,channel: U32,input: Bands.Bands(~Element),planned: Bands.Bands(~Element) & Bands.Plan) -> Bands.Bands(~Element):
128  (output,+layout) = planned
129  Bands.each(~Bands.Paired<Storage.Buffer<Element>,Storage.Buffer<Element>>,~Storage.Buffer<Element>,~Unit,~Channel,~Bands.unshared,
130    ~Bands.paired(~Storage.Buffer<Element>,~Storage.Buffer<Element>,~Unit,~Channel,~append_band(~Element)),
131    ~Bands.left_of(~Storage.Buffer<Element>,~Storage.Buffer<Element>),
132    layout,parallel && Bands.same_image(geometry,layout),parallel,geometry,Channel{geometry,channel},Bands.zip(~Storage.Buffer<Element>,~Storage.Buffer<Element>,output,input),Unit{})
133
134# Every output leaf receives the same rows of the input, through the shared
135# band recursion on the plan of the output. Trees of another shape keep the output.
136def append_aligned(~Element: Data,parallel: Bool,geometry: Bands.Geometry,channel: U32,input: Bands.Bands(~Element),output: Bands.Bands(~Element)) -> Bands.Bands(~Element):
137  append_planned(~Element,parallel,geometry,channel,input,Bands.plan(~Element,output))
138
139def appended_rows(~Element: Data,owners: Tensors.Tensor<Element> & Bands.Bands(~Element)) -> Bands.Bands(~Element):
140  (source,output) = owners
141  output
142
143def selected_rows(~Element: Data,same: Bool,+geometry: Bands.Geometry,channel: U32,workers: U32,layout: Bands.Plan,input: Bands.Bands(~Element),output: Bands.Bands(~Element),shape: List<&2,U32>,fill: Element,status: Tensors.Status,previous: U32) -> Bands.Bands(~Element):
144  match same:
145    case True{}: append_aligned(~Element,BandSchedule.parallel(geometry,workers,8.0,0.0,layout),geometry,channel,input,output)
146    case False{}:
147      appended_rows(~Element,append_tree(~Element,output,geometry,channel,Tensors.Banded{input,geometry,shape,fill,status,previous}))
148
149def aligned_or_rows(~Element: Data,+geometry: Bands.Geometry,channel: U32,workers: U32,
150  input: Bands.Bands(~Element) & Bands.Plan,output: Bands.Bands(~Element) & Bands.Plan,shape: List<&2,U32>,fill: Element,status: Tensors.Status,previous: U32) -> Bands.Bands(~Element):
151  (input,source_plan) = input
152  (output,+target_plan) = output
153  selected_rows(~Element,Bands.same_plan(source_plan,target_plan),geometry,channel,workers,target_plan,input,output,shape,fill,status,previous)
154
155def append_storage(~Element: Data,+geometry: Bands.Geometry,channel: U32,workers: U32,input: Tensors.Tensor<Element>,output: Bands.Bands(~Element)) -> Bands.Bands(~Element):
156  match input:
157    case Tensors.Tensor{storage,shape,fill,status,previous}:
158      appended_rows(~Element,append_tree(~Element,output,geometry,channel,Tensors.Tensor{storage,shape,fill,status,previous}))
159    case Tensors.Banded{bands,layout,shape,fill,status,previous}:
160      aligned_or_rows(~Element,geometry,channel,workers,Bands.plan(~Element,bands),Bands.plan(~Element,output),shape,fill,status,previous)
161
162def append_input(~Element: Data,geometry: Bands.Geometry,+channel: U32,+count: U32,workers: U32,input: Tensors.Tensor<Element>,output: Bands.Bands(~Element)) -> Bands.Bands(~Element) & U32:
163  Bands.Geometry{batch,channels,height,width} = geometry
164  (append_storage(~Element,Bands.Geometry{batch,count,height,width},channel,workers,input,output),(channel + count : U32))
165
166def append_part(~Element: Data,geometry: Bands.Geometry,workers: U32,input: Tensors.Tensor<Element>,state: Bands.Bands(~Element) & U32) -> Bands.Bands(~Element) & U32:
167  (output,channel) = state
168  Tensors.with_metadata(~Element,~(Bands.Bands(~Element) & U32),input,input => shape => fill => status => previous =>
169    append_input(~Element,geometry,channel,channels(shape),workers,input,output))
170
171def output(~Element: Data,state: Bands.Bands(~Element) & U32) -> Bands.Bands(~Element):
172  (output,channel) = state
173  output
174
175def append_parts(~Element: Data,parts: Tensors.Parts<Element>,+geometry: Bands.Geometry,+workers: U32,state: Bands.Bands(~Element) & U32) -> Bands.Bands(~Element):
176  match parts:
177    case Tensors.Single{input}: output(~Element,append_part(~Element,geometry,workers,input,state))
178    case Tensors.Cons{input,rest}: append_parts(~Element,rest,geometry,workers,append_part(~Element,geometry,workers,input,state))
179
180def join(~Element: Data,layout: Bands.Plan,+geometry: Bands.Geometry,parts: Tensors.Parts<Element>,shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
181  Tensors.Banded{append_parts(~Element,parts,geometry,workers,(allocate(~Element,layout,geometry,fill),0)),geometry,shape,fill,Tensors.Usable{},workers}