tensor_access.bend
stelliferous@0.0.3.0 · 213 lines · 0 definitions
1# Internal tensor read layouts and unchecked traversal kernels.
2import Base
3import stelliferous@0.0.3.0/tensor_view.bend as Views
4import stelliferous@0.0.3.0/traversal_array.bend as Traversal
5import stelliferous@0.0.3.0/traversal_loop.bend as Loop
6import stelliferous@0.0.3.0/tensor_cursor.bend as Cursor
7
8# A strided operation visits the source view in row-major order. It proceeds
9# in runs of the innermost axis: every element of a run but the last advances
10# only the address, and the last takes the ordinary cursor step, which carries
11# into the outer axes. The element loop is a small native loop whatever else
12# shares the cursor; stepping the whole cursor per element costs several calls
13# with a wide state once more than one operation uses it. inside and last are
14# the same element action, stepping a stride and the whole cursor.
15def last_of_run(~Element: Data,extent: Nat,stride: U32,rewind: U32,rest: Cursor.Depth1(),
16 state: Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Stride>) -> Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Cursor>:
17 Traversal.State{input,output,position,Cursor.Stride{address,step}} = state
18 Traversal.State{input,output,position,Cursor.Cursor{Cursor.Level{Cursor.Axis{extent,1n,stride,rewind},rest},address}}
19
20def run_axis(~Element: Data,~inside: Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Stride> -> Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Stride>,~last: Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Cursor> -> Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Cursor>,remaining: Nat,extent: Nat,+stride: U32,rewind: U32,rest: Cursor.Depth1(),+address: U32,
21 input: Array<Element>,output: Array<Element>,position: U32) -> Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Cursor>:
22 match remaining:
23 case 0n:
24 last(Traversal.State{input,output,position,Cursor.Cursor{Cursor.Level{Cursor.Axis{extent,0n,stride,rewind},rest},address}})
25 case 1n+count:
26 last(last_of_run(~Element,extent,stride,rewind,rest,
27 Loop.run(~Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Stride>,~inside,count,Traversal.State{input,output,position,Cursor.Stride{address,stride}})))
28
29def run_step(~Element: Data,~inside: Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Stride> -> Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Stride>,~last: Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Cursor> -> Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Cursor>,state: Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Cursor>) ->
30 Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Cursor>:
31 Traversal.State{input,output,position,Cursor.Cursor{Cursor.Level{Cursor.Axis{extent,remaining,stride,rewind},rest},address}} = state
32 run_axis(~Element,~inside,~last,remaining,extent,stride,rewind,rest,address,input,output,position)
33
34# Runs are used when the cursor starts a full run and count is a whole number
35# of runs; any other call steps element by element. Both visit the same
36# elements in the same order.
37def run_guarded(~Element: Data,~inside: Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Stride> -> Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Stride>,~last: Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Cursor> -> Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Cursor>,full_run: Bool,whole_runs: Bool,runs: Nat,count: Nat,cursor: Cursor.Cursor,destination: U32,
38 input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:
39 match full_run whole_runs:
40 case True{} True{}:
41 Traversal.finish(~Array<Element>,~Element,~Cursor.Cursor,
42 Loop.run(~Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Cursor>,~run_step(~Element,~inside,~last),runs,
43 Traversal.State{input,output,destination,cursor}))
44 case full_run whole_runs:
45 Traversal.finish(~Array<Element>,~Element,~Cursor.Cursor,
46 Loop.run(~Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Cursor>,~last,count,Traversal.State{input,output,destination,cursor}))
47
48def run_remaining(~Element: Data,~inside: Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Stride> -> Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Stride>,~last: Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Cursor> -> Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Cursor>,remaining: Nat,+extent: Nat,stride: U32,rewind: U32,rest: Cursor.Depth1(),address: U32,+count: Nat,destination: U32,
49 input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:
50 match remaining:
51 case 0n:
52 Traversal.finish(~Array<Element>,~Element,~Cursor.Cursor,
53 Loop.run(~Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Cursor>,~last,count,
54 Traversal.State{input,output,destination,Cursor.Cursor{Cursor.Level{Cursor.Axis{extent,0n,stride,rewind},rest},address}}))
55 case 1n+ +inner_count:
56 +runs = Nat.div(count,1n+inner_count)
57 run_guarded(~Element,~inside,~last,Nat.is_eq(extent,1n+inner_count),Nat.is_eq(Nat.mul(runs,1n+inner_count),count),runs,count,
58 Cursor.Cursor{Cursor.Level{Cursor.Axis{extent,1n+inner_count,stride,rewind},rest},address},destination,input,output)
59
60def run_strided(~Element: Data,~inside: Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Stride> -> Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Stride>,~last: Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Cursor> -> Traversal.State<&1,&1,Array<Element>,Array<Element>,Cursor.Cursor>,count: Nat,cursor: Cursor.Cursor,destination: U32,input: Array<Element>,output: Array<Element>) ->
61 Array<Element> & Array<Element>:
62 Cursor.Cursor{Cursor.Level{Cursor.Axis{extent,remaining,stride,rewind},rest},address} = cursor
63 run_remaining(~Element,~inside,~last,remaining,extent,stride,rewind,rest,address,count,destination,input,output)
64
65def copy_strided(~Element: Data,count: Nat,cursor: Cursor.Cursor,destination: U32,input: Array<Element>,output: Array<Element>) ->
66 Array<Element> & Array<Element>:
67 run_strided(~Element,~Traversal.visit_step(~Array<Element>,~Element,~Cursor.Stride,~Cursor.read_stride(~Element)),
68 ~Traversal.visit_step(~Array<Element>,~Element,~Cursor.Cursor,~Cursor.read(~Element)),count,cursor,destination,input,output)
69
70# output[destination + i] = operation(output[destination + i], element i of the view).
71def update_strided(~Element: Data,~operation: Element -> Element -> Element,count: Nat,cursor: Cursor.Cursor,destination: U32,
72 input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:
73 run_strided(~Element,~Traversal.visit_update_step(~Array<Element>,~Element,~Cursor.Stride,~Cursor.read_stride(~Element),~operation),
74 ~Traversal.visit_update_step(~Array<Element>,~Element,~Cursor.Cursor,~Cursor.read(~Element),~operation),count,cursor,destination,input,output)
75
76# output[address of element i of the cursor] = input[source + i].
77def scatter_strided(~Element: Data,count: Nat,cursor: Cursor.Cursor,source: U32,input: Array<Element>,output: Array<Element>) ->
78 Array<Element> & Array<Element>:
79 run_strided(~Element,~Traversal.spread_step(~Element,~Cursor.Stride,~Cursor.target_stride),
80 ~Traversal.spread_step(~Element,~Cursor.Cursor,~Cursor.target),count,cursor,source,input,output)
81
82def scatter_run(~Element: Data,count: Nat,target: Views.View,source: U32,input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:
83 scatter_strided(~Element,count,Cursor.prepare(target),source,input,output)
84
85def copy_run(~Element: Data,contiguous: Bool,count: Nat,source: Views.View,destination: U32,
86 input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:
87 match contiguous:
88 case True{}:
89 Views.View{axes,offset} = source
90 Traversal.copy(~Element,count,offset,destination,input,output)
91 case False{}: copy_strided(~Element,count,Cursor.prepare(source),destination,input,output)
92
93def zero_contiguous(+view: Views.View) -> Bool:
94 Views.View{axes,offset} = view
95 U32.is_zero(offset) && Views.is_contiguous(view)
96
97def combine_run(~Element: Data,~operation: Element -> Element -> Element,fast: Bool,count: Nat,
98 right_view: Views.View,destination: U32,right: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:
99 match fast:
100 case True{}:
101 Traversal.finish(~Array<Element>,~Element,~Unit,
102 Traversal.update(~Array<Element>,~Element,~Element,~Unit,~Traversal.read_index(~Element),~operation,~Traversal.index,count,
103 Traversal.State{right,output,0,Unit{}}))
104 case False{}: update_strided(~Element,~operation,count,Cursor.prepare(right_view),destination,right,output)
105
106def apply_element(~Element: Data,~operation: Element -> Element,context: Unit,value: Element) -> Element:
107 operation(value)
108
109# Reduction of the lines of a view, one output at a time. The cursor walks
110# the view without its last axis, which has extent values stride apart; a
111# line is runs consecutive runs of that axis. Output i folds the values of
112# line i in the view's row-major order, from initial, and stores finish of
113# the fold.
114type Runs is Data:
115 Runs{cursor: Cursor.Cursor,extent: U32,stride: U32}
116
117type Lines<-Accumulator: Data> is Data:
118 Lines{cursor: Cursor.Cursor,runs: Nat,extent: U32,stride: U32,initial: Accumulator}
119
120def ran(~Element: Data,~Accumulator: Data,axes: Cursor.Levels(),address: U32,+extent: U32,+stride: U32,+position: U32,
121 result: Array<Element> & Accumulator) -> Traversal.State<&1,&2,Array<Element>,Accumulator,Runs>:
122 (input,folded) = result
123 Traversal.State{input,folded,(position + 1 : U32),Runs{Cursor.advance(axes,address),extent,stride}}
124
125# One run folded into the accumulator; the cursor moves to the next run.
126def fold_run(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,
127 state: Traversal.State<&1,&2,Array<Element>,Accumulator,Runs>) -> Traversal.State<&1,&2,Array<Element>,Accumulator,Runs>:
128 Traversal.State{input,folded,+position,Runs{Cursor.Cursor{axes,+address},+extent,+stride}} = state
129 ran(~Element,~Accumulator,axes,address,extent,stride,position,
130 Traversal.fold_rounds(~Array<Element>,~Element,~Accumulator,~Traversal.Strided,~Traversal.strided(~Element),~operation,U32.to_nat(extent),0,input,folded,
131 Traversal.Strided{address,stride}))
132
133def folded_line(~Element: Data,~Accumulator: Data,~finish: Accumulator -> Element,+runs: Nat,+initial: Accumulator,
134 state: Traversal.State<&1,&2,Array<Element>,Accumulator,Runs>) -> Array<Element> & Traversal.Visited<Lines<Accumulator>,Element>:
135 Traversal.State{input,folded,position,Runs{cursor,extent,stride}} = state
136 (input,Traversal.Visited{Lines{cursor,runs,extent,stride,initial},finish(folded)})
137
138def fold_line(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element,
139 input: Array<Element>,context: Lines<Accumulator>) -> Array<Element> & Traversal.Visited<Lines<Accumulator>,Element>:
140 Lines{cursor,+runs,extent,stride,+initial} = context
141 folded_line(~Element,~Accumulator,~finish,runs,initial,
142 Loop.run(~Traversal.State<&1,&2,Array<Element>,Accumulator,Runs>,~fold_run(~Element,~Accumulator,~operation),runs,
143 Traversal.State{input,initial,0,Runs{cursor,extent,stride}}))
144
145# count lines of runs runs; view is the view without its last axis.
146def reduce_lines(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element,
147 count: Nat,view: Views.View,+runs: Nat,+extent: U32,+stride: U32,+initial: Accumulator,input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:
148 Traversal.finish(~Array<Element>,~Element,~Lines<Accumulator>,
149 Traversal.visit(~Array<Element>,~Element,~Lines<Accumulator>,~fold_line(~Element,~Accumulator,~operation,~finish),count,
150 Traversal.State{input,output,0,Lines{Cursor.prepare(view),runs,extent,stride,initial}}))
151
152# The geometry of a two-dimensional window reduction over the last two axes.
153# It runs as two sweeps: along the columns of every row, then along the rows.
154type Pool is Data:
155 Pool{height: U32,width: U32,kernel: U32,stride: U32,padding: U32,out_height: U32,out_width: U32}
156
157def window_read(~Element: Data,+pad: Element,valid: Bool,input: Array<Element>,index: U32) -> Array<Element> & Element:
158 match valid:
159 case False{}: (input,pad)
160 case True{}: Array.get(Element,input,index)
161
162# Window reduction along one axis. step is the number of elements after the
163# axis, the same in the input and the output. Output position p keeps its
164# step offset p % step, lies at coordinate (p / step) % out_extent of the
165# axis, and folds kernel samples in increasing coordinate from pad; a sample
166# outside the axis reads pad. origin is the coordinate at which the window of
167# output coordinate zero starts before padding: zero for a whole axis, the
168# stored row for the rows of a taller image.
169type Sweep is Data:
170 Sweep{extent: U32,step: U32,kernel: U32,stride: U32,padding: U32,out_extent: U32,origin: U32}
171
172type SweepContext<-Element: Data> is Data:
173 SweepContext{sweep: Sweep,pad: Element}
174
175type Line<-Element: Data> is Data:
176 Line{base: U32,first: U32,extent: U32,step: U32,padding: U32,pad: Element}
177
178def read_line(~Element: Data,input: Array<Element>,line: Line<Element>,position: U32) -> Array<Element> & Element:
179 Line{+base,+first,+extent,+step,+padding,+pad} = line
180 +coordinate = (first + position - padding : U32)
181 window_read(~Element,pad,U32.is_lt(coordinate,extent),input,(base + coordinate * step : U32))
182
183def line_at(~Element: Data,+sweep: Sweep,+pad: Element,+position: U32) -> Line<Element>:
184 Sweep{+extent,+step,kernel,+stride,+padding,+out_extent,+origin} = sweep
185 +along = (position / step : U32)
186 Line{(along / out_extent * extent * step + position % step : U32),(along % out_extent * stride + origin : U32),extent,step,padding,pad}
187
188def read_swept(~Element: Data,~operation: Element -> Element -> Element,input: Array<Element>,context: SweepContext<Element>,position: U32) ->
189 Array<Element> & Element:
190 SweepContext{+sweep,+pad} = context
191 Sweep{extent,step,+kernel,stride,padding,out_extent,origin} = sweep
192 Traversal.fold(~Array<Element>,~Element,~Element,~Line<Element>,~read_line(~Element),~operation,U32.to_nat(kernel),0,input,pad,
193 line_at(~Element,sweep,pad,position))
194
195def swept_index(~Element: Data,context: SweepContext<Element>,position: U32) -> U32:
196 position
197
198def sweep_run(~Element: Data,~operation: Element -> Element -> Element,count: Nat,context: SweepContext<Element>,
199 input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:
200 Traversal.finish(~Array<Element>,~Element,~SweepContext<Element>,
201 Traversal.run(~Array<Element>,~Element,~SweepContext<Element>,~read_swept(~Element,~operation),~swept_index(~Element),count,
202 Traversal.State{input,output,0,context}))
203
204# The sweep along the columns of a pool: every stored row is one line.
205def pool_columns(pool: Pool) -> Sweep:
206 Pool{height,width,kernel,stride,padding,out_height,out_width} = pool
207 Sweep{width,1,kernel,stride,padding,out_width,0}
208
209# The sweep along the rows of the column-swept planes. The stored rows may be
210# a window of a taller image; see Sweep for origin.
211def pool_rows(pool: Pool,origin: U32) -> Sweep:
212 Pool{height,width,kernel,stride,padding,out_height,out_width} = pool
213 Sweep{height,out_width,kernel,stride,padding,out_height,origin}