1 # A reduction runs on the lines of a view of its input: the view puts the 2 # reduced axes last, and line k holds their values at index k of the kept 3 # axes. The view is classified over the input's bands as a gather is: when 4 # the bands keep their images and rows and the reduced axes lie after the 5 # rows, every band folds the lines of its own rows through its local view; 6 # otherwise the input arrives as one band, which folds every line. 7 import Base 8 import ./storage_buffer.bend as Storage 9 import ./tensor_view.bend as Views 10 import ./tensor_operations.bend as Operations 11 import ./tensor_band.bend as Bands 12 import ./tensor_band_view.bend as BandViews 13 import ./traversal_partition.bend as Partition 14 import ./tensor.bend as Tensors 15 import ./tensor_band_schedule.bend as BandSchedule 16 17 # What every band of a reduction shares: the banded view of the input 18 # geometry, the number of reduced axes, the initial accumulator and the fill 19 # of fresh storage. 20 type Context<-Accumulator: Data,-Element: Data> is Data: 21 Context{banded: BandViews.Banded,reduced: Nat,geometry: Bands.Geometry,initial: Accumulator,fill: Element} 22 23 # A band's folds as a checked band: its outputs beside whether the call 24 # accepted. 25 def checked(~Element: Data,+first: U32,+rows: U32,result: Operations.Outcome<(Storage.Buffer & Storage.Buffer)>) -> Bands.Transformed(~Element): 26 match result: 27 case Operations.Accepted{(input,output)}: Partition.Leaf{Bands.Checked{output,True{}},first,rows,Unit{}} 28 case Operations.Rejected{(input,output)}: Partition.Leaf{Bands.Checked{output,False{}},first,rows,Unit{}} 29 30 # The axes of a view before its last reduced axes. 31 def kept_axes(+axes: List<&2,Views.Axis>,+reduced: Nat) -> List<&2,Views.Axis>: 32 List.take(&2,Views.Axis,axes,Nat.sub(List.length(&2,Views.Axis,axes),reduced)) 33 34 def folded_view(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element, 35 +reduced: Nat,+initial: Accumulator,+fill: Element,+first: U32,+rows: U32,local: Views.View,input: Storage.Buffer) -> Bands.Transformed(~Element): 36 Views.View{+axes,+offset} = local 37 outputs = {Views.View{kept_axes(axes,reduced),0} : Views.View} 38 checked(~Element,first,rows, 39 Operations.reduce_lines(~Element,~Accumulator,~operation,~finish,input,axes,reduced,offset,initial, 40 Tensors.buffer(~Element,fill,BandViews.known(Views.numel(outputs))))) 41 42 # The band of rows [first, first + rows) folds the lines of its rows; its 43 # output band holds their folds. A band of no rows has no line. 44 def folded_rows(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element, 45 empty: Bool,+banded: BandViews.Banded,+reduced: Nat,+geometry: Bands.Geometry,+initial: Accumulator,+fill: Element,+first: U32,+rows: U32, 46 input: Storage.Buffer) -> Bands.Transformed(~Element): 47 match empty: 48 case True{}: 49 +scale = BandViews.scale(banded) 50 Partition.Leaf{Bands.Checked{Tensors.buffer(~Element,fill,0),True{}},(first * scale : U32),(rows * scale : U32),Unit{}} 51 case False{}: 52 +scale = BandViews.scale(banded) 53 folded_view(~Element,~Accumulator,~operation,~finish,reduced,initial,fill,(first * scale : U32),(rows * scale : U32), 54 BandViews.local(banded,geometry,rows),input) 55 56 def leaf(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element, 57 context: Context,first: U32,rows: U32,input: Storage.Buffer) -> Bands.Transformed(~Element): 58 Context{+banded,+reduced,+geometry,+initial,+fill} = context 59 +first = first 60 +rows = rows 61 folded_rows(~Element,~Accumulator,~operation,~finish,U32.is_zero(rows),banded,reduced,geometry,initial,fill,first,rows,input) 62 63 # The geometry of the outputs: the same images, the axes before the rows as 64 # channels, each row repeated, the kept axes after the rows as columns. 65 def output(+banded: BandViews.Banded,+geometry: Bands.Geometry,+reduced: Nat) -> Bands.Geometry: 66 BandViews.Banded{before,repeats,after,offset} = banded 67 Bands.Geometry{batch,channels,height,width} = geometry 68 Bands.Geometry{batch,BandViews.product_or_zero(Views.axis_extents(before)),(height * BandViews.product_or_zero(repeats) : U32), 69 BandViews.product_or_zero(Views.axis_extents(kept_axes(after,reduced)))} 70 71 # The banded view when the reduced axes lie among the axes after the rows. 72 def in_band(found: Maybe<&2,BandViews.Banded>,+reduced: Nat) -> Maybe<&2,BandViews.Banded>: 73 match found: 74 case Some{BandViews.Banded{+before,+repeats,+after,+offset}}: 75 BandViews.kept_when(Nat.is_le(reduced,List.length(&2,Views.Axis,after)),BandViews.Banded{before,repeats,after,offset}) 76 case None{}: None{} 77 78 # Every band folds its own lines. operations prices the combination of one 79 # value, beside the read of it. 80 def banded_lines(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element,+name: String, 81 +operations: F32,+initial: Accumulator,+reduced: Nat,+banded: BandViews.Banded,+geometry: Bands.Geometry,+shape: List<&2,U32>,+fill: Element,+workers: U32, 82 bands: Bands.Bands(~Element)) -> Tensors.Tensor: 83 Tensors.transformed(~Element,name,Tensors.InvalidArgument{shape},output(banded,geometry,reduced),shape,fill,workers, 84 BandSchedule.transform(~Element,~Context,~leaf(~Element,~Accumulator,~operation,~finish),Context{banded,reduced,geometry,initial,fill}, 85 geometry,workers,4.0,operations,bands)) 86 87 def whole_lines(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element,+name: String, 88 +operations: F32,+initial: Accumulator,+reduced: Nat,found: Maybe<&2,BandViews.Banded>,+whole: Bands.Geometry,+shape: List<&2,U32>,+fill: Element,+workers: U32, 89 tensor: Tensors.Tensor) -> Tensors.Tensor: 90 match found: 91 case Some{+banded}: 92 banded_lines(~Element,~Accumulator,~operation,~finish,name,operations,initial,reduced,banded,whole,shape,fill,workers, 93 Tensors.arrive(~Element,whole,Bands.unsplit(whole),tensor)) 94 case None{}: Tensors.failed(~Element,name,Tensors.InvalidArgument{shape},shape,fill,workers) 95 96 def found_lines(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element,+name: String, 97 +operations: F32,+initial: Accumulator,+reduced: Nat,found: Maybe<&2,BandViews.Banded>,+axes: List<&2,Views.Axis>,+offset: U32,bands: Bands.Bands(~Element), 98 +geometry: Bands.Geometry,+source_shape: List<&2,U32>,+source_fill: Element,+source_workers: U32,+shape: List<&2,U32>,+fill: Element,+workers: U32) -> 99 Tensors.Tensor: 100 match found: 101 case Some{+banded}: banded_lines(~Element,~Accumulator,~operation,~finish,name,operations,initial,reduced,banded,geometry,shape,fill,workers,bands) 102 case None{}: 103 +whole = {Bands.Geometry{1,1,1,Tensors.count_or_zero(source_shape)} : Bands.Geometry} 104 whole_lines(~Element,~Accumulator,~operation,~finish,name,operations,initial,reduced,in_band(BandViews.classify(axes,offset,whole),reduced),whole,shape,fill,workers, 105 Tensors.Tensor{bands,geometry,source_shape,source_fill,Tensors.Usable{},source_workers}) 106 107 # No output, or lines of no value, leave every output at finish(initial). 108 def lines_when(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element,+name: String, 109 empty: Bool,+operations: F32,+count: U32,+initial: Accumulator,+reduced: Nat,+axes: List<&2,Views.Axis>,+offset: U32,bands: Bands.Bands(~Element), 110 +geometry: Bands.Geometry,+source_shape: List<&2,U32>,+source_fill: Element,+source_workers: U32,+shape: List<&2,U32>,+fill: Element,+workers: U32) -> 111 Tensors.Tensor: 112 match empty: 113 case True{}: Tensors.whole(~Element,Tensors.buffer(~Element,finish(initial),count),shape,fill,Tensors.Usable{},workers) 114 case False{}: 115 found_lines(~Element,~Accumulator,~operation,~finish,name,operations,initial,reduced,in_band(BandViews.classify(axes,offset,geometry),reduced), 116 axes,offset,bands,geometry,source_shape,source_fill,source_workers,shape,fill,workers) 117 118 # count outputs of shape: output k folds the length values of line k of the 119 # view, the values of its last reduced axes at index k of the axes before 120 # them, in row-major order from initial, and holds finish of the fold. A 121 # failed tensor keeps its problem. 122 def lines(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element,+name: String, 123 +operations: F32,+count: U32,+length: U32,+initial: Accumulator,+shape: List<&2,U32>,+fill: Element,+workers: U32,tensor: Tensors.Tensor, 124 view: Maybe<&2,Views.View>,+reduced: Nat) -> Tensors.Tensor: 125 Tensors.Tensor{bands,+geometry,+source_shape,+source_fill,status,+source_workers} = tensor 126 match status view: 127 case Tensors.Usable{} Some{Views.View{+axes,+offset}}: 128 lines_when(~Element,~Accumulator,~operation,~finish,name,U32.is_zero(count) || U32.is_zero(length),operations,count,initial,reduced,axes,offset,bands, 129 geometry,source_shape,source_fill,source_workers,shape,fill,workers) 130 case Tensors.Usable{} None{}: Tensors.failed(~Element,name,Tensors.InvalidArgument{source_shape},shape,fill,workers) 131 case Tensors.Failed{problem} view: Tensors.whole(~Element,Tensors.buffer(~Element,fill,0),shape,fill,Tensors.Failed{problem},workers)