1 # Channel concatenation writes each input range directly into its final band. 2 # Different input boundaries need no temporary dense tensor or aligned copy. 3 import Base 4 import ./tensor.bend as Tensors 5 import ./tensor_band.bend as Bands 6 import ./storage_buffer.bend as Storage 7 import ./traversal_partition.bend as Partition 8 import ./storage_copy.bend as Copy 9 import ./tensor_band_schedule.bend as BandSchedule 10 11 def 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 17 def 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 & Maybe<&2,Bands.Plan>: 19 (bands,layout) = owners 20 (Tensors.Banded{bands,geometry,shape,fill,status,workers},Some{layout}) 21 22 def tensor_plan(~Element: Data,tensor: Tensors.Tensor) -> Tensors.Tensor & 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 27 def 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 32 def single_plan(~Element: Data,first: Tensors.Tensor & Maybe<&2,Bands.Plan>) -> Tensors.Parts & Maybe<&2,Bands.Plan>: 33 (tensor,layout) = first 34 (Tensors.Single{tensor},layout) 35 36 def cons_plan(~Element: Data,first: Tensors.Tensor & Maybe<&2,Bands.Plan>,rest: Tensors.Parts & Maybe<&2,Bands.Plan>) -> Tensors.Parts & Maybe<&2,Bands.Plan>: 37 (tensor,first) = first 38 (rest,other) = rest 39 (Tensors.Cons{tensor,rest},first_plan(first,other)) 40 41 def discover(~Element: Data,parts: Tensors.Parts) -> Tensors.Parts & 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. 48 def 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 53 def single_compatible_plan(~Element: Data,tensor: Tensors.Tensor,layout: Maybe<&2,Bands.Plan>) -> 54 (Tensors.Parts & 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 59 def single_compatible(~Element: Data,first: Tensors.Tensor & Maybe<&2,Bands.Plan>) -> 60 (Tensors.Parts & Maybe<&2,Bands.Plan>) & Bool: 61 (tensor,layout) = first 62 single_compatible_plan(~Element,tensor,layout) 63 64 def cons_compatible(~Element: Data,first: Tensors.Tensor & Maybe<&2,Bands.Plan>, 65 rest: (Tensors.Parts & Maybe<&2,Bands.Plan>) & Bool) -> (Tensors.Parts & 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 73 def discover_compatible(~Element: Data,parts: Tensors.Parts) -> (Tensors.Parts & 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 78 def 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 83 def dense_rows(~Element: Data,geometry: Bands.Geometry,+first: U32,+rows: U32,channel: U32,input: Storage.Buffer,output: Storage.Buffer) -> Storage.Buffer & Storage.Buffer: 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 88 def channel_offset(geometry: Bands.Geometry,channel: U32,rows: U32) -> U32: 89 Bands.Geometry{batch,channels,height,width} = geometry 90 (channel * rows * width : U32) 91 92 def rows(~Element: Data,geometry: Bands.Geometry,+first: U32,+rows: U32,channel: U32,tensor: Tensors.Tensor,output: Storage.Buffer) -> Tensors.Tensor & Storage.Buffer: 93 match tensor: 94 case Tensors.Tensor{input,shape,fill,status,workers}: 95 Storage.bind_pair(~Storage.Buffer,~Storage.Buffer,~(Tensors.Tensor & Storage.Buffer), 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,~(Tensors.Tensor & Storage.Buffer), 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 102 def append_tree(~Element: Data,output: Bands.Bands(~Element),+geometry: Bands.Geometry,+channel: U32,input: Tensors.Tensor) -> Tensors.Tensor & Bands.Bands(~Element): 103 match output: 104 case Partition.Leaf{output,+first,+count,details}: 105 Storage.bind_pair(~Tensors.Tensor,~Storage.Buffer,~(Tensors.Tensor & 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,~Bands.Bands(~Element),~(Tensors.Tensor & Bands.Bands(~Element)), 109 append_tree(~Element,left,geometry,channel,input),input => left => 110 Storage.bind_pair(~Tensors.Tensor,~Bands.Bands(~Element),~(Tensors.Tensor & Bands.Bands(~Element)), 111 append_tree(~Element,right,geometry,channel,input),input => right => (input,Partition.Fork{left,right}))) 112 113 def 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 119 type Channel is Data: 120 Channel{geometry: Bands.Geometry,channel: U32} 121 122 def append_band(~Element: Data,context: Channel,first: U32,rows: U32,output: Storage.Buffer,input: Storage.Buffer,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 127 def 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>,~Storage.Buffer,~Unit,~Channel,~Bands.unshared, 130 ~Bands.paired(~Storage.Buffer,~Storage.Buffer,~Unit,~Channel,~append_band(~Element)), 131 ~Bands.left_of(~Storage.Buffer,~Storage.Buffer), 132 layout,parallel && Bands.same_image(geometry,layout),parallel,geometry,Channel{geometry,channel},Bands.zip(~Storage.Buffer,~Storage.Buffer,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. 136 def 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 139 def appended_rows(~Element: Data,owners: Tensors.Tensor & Bands.Bands(~Element)) -> Bands.Bands(~Element): 140 (source,output) = owners 141 output 142 143 def 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 149 def 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 155 def append_storage(~Element: Data,+geometry: Bands.Geometry,channel: U32,workers: U32,input: Tensors.Tensor,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 162 def append_input(~Element: Data,geometry: Bands.Geometry,+channel: U32,+count: U32,workers: U32,input: Tensors.Tensor,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 166 def append_part(~Element: Data,geometry: Bands.Geometry,workers: U32,input: Tensors.Tensor,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 171 def output(~Element: Data,state: Bands.Bands(~Element) & U32) -> Bands.Bands(~Element): 172 (output,channel) = state 173 output 174 175 def append_parts(~Element: Data,parts: Tensors.Parts,+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 180 def join(~Element: Data,layout: Bands.Plan,+geometry: Bands.Geometry,parts: Tensors.Parts,shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensors.Tensor: 181 Tensors.Banded{append_parts(~Element,parts,geometry,workers,(allocate(~Element,layout,geometry,fill),0)),geometry,shape,fill,Tensors.Usable{},workers}