tensor_band_reduce.bend
stelliferous@0.0.3.0 · 131 lines · 0 definitions
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.
7import Base
8import stelliferous@0.0.3.0/storage_buffer.bend as Storage
9import stelliferous@0.0.3.0/tensor_view.bend as Views
10import stelliferous@0.0.3.0/tensor_operations.bend as Operations
11import stelliferous@0.0.3.0/tensor_band.bend as Bands
12import stelliferous@0.0.3.0/tensor_band_view.bend as BandViews
13import stelliferous@0.0.3.0/traversal_partition.bend as Partition
14import stelliferous@0.0.3.0/tensor.bend as Tensors
15import stelliferous@0.0.3.0/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.
20type 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.
25def checked(~Element: Data,+first: U32,+rows: U32,result: Operations.Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>) -> 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.
31def 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
34def 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<Element>) -> 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.
44def 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<Element>) -> 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
56def leaf(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element,
57 context: Context<Accumulator,Element>,first: U32,rows: U32,input: Storage.Buffer<Element>) -> 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.
65def 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.
72def 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.
80def 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<Element>:
83 Tensors.transformed(~Element,name,Tensors.InvalidArgument{shape},output(banded,geometry,reduced),shape,fill,workers,
84 BandSchedule.transform(~Element,~Context<Accumulator,Element>,~leaf(~Element,~Accumulator,~operation,~finish),Context{banded,reduced,geometry,initial,fill},
85 geometry,workers,4.0,operations,bands))
86
87def 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<Element>) -> Tensors.Tensor<Element>:
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
96def 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<Element>:
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).
108def 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<Element>:
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.
122def 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<Element>,
124 view: Maybe<&2,Views.View>,+reduced: Nat) -> Tensors.Tensor<Element>:
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)