tensor_access.bend

stelliferous@0.0.2.0 · 199 lines · 30 definitions

1# Internal tensor read layouts and unchecked traversal kernels.
2import Base
3import stelliferous@0.0.2.0/tensor_view.bend as Views
4import stelliferous@0.0.2.0/traversal_array.bend as Traversal
5import stelliferous@0.0.2.0/traversal_loop.bend as Loop
6import stelliferous@0.0.2.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
109def reduce_run(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,
110  fast: Bool,count: Nat,view: Views.View,storage: Array<Element>,initial: Accumulator) -> Array<Element> & Accumulator:
111  match fast:
112    case True{}:
113      Views.View{axes,offset} = view
114      Traversal.fold(~Array<Element>,~Element,~Accumulator,~Traversal.Offsets,~Traversal.source_offset(~Element),~operation,
115        count,0,storage,initial,Traversal.Offsets{offset,0})
116    case False{}:
117      Traversal.finish_fold(~Array<Element>,~Accumulator,~Cursor.Cursor,
118        Traversal.visit_fold(~Array<Element>,~Element,~Accumulator,~Cursor.Cursor,~Cursor.read(~Element),~operation,count,
119          Traversal.State{storage,initial,0,Cursor.prepare(view)}))
120
121def indexed_value(~Element: Data,~operation: U32 -> Element -> Element,+position: U32,result: Array<Element> & Element) -> Array<Element> & Element:
122  (input,value) = result
123  (input,operation(position,value))
124
125def read_indexed(~Element: Data,~operation: U32 -> Element -> Element,input: Array<Element>,context: Traversal.Offsets,position: U32) ->
126  Array<Element> & Element:
127  +position = position
128  indexed_value(~Element,~operation,position,Traversal.source_offset(~Element,input,context,position))
129
130# Reads a contiguous view and writes a separate contiguous output from zero.
131def apply_indexed_run(~Element: Data,~operation: U32 -> Element -> Element,count: Nat,view: Views.View,
132  input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:
133  Views.View{axes,offset} = view
134  Traversal.finish(~Array<Element>,~Element,~Traversal.Offsets,
135    Traversal.run(~Array<Element>,~Element,~Traversal.Offsets,~read_indexed(~Element,~operation),~Traversal.destination_offset,count,
136      Traversal.State{input,output,0,Traversal.Offsets{offset,0}}))
137
138# The geometry of a two-dimensional window reduction over the last two axes.
139# It runs as two sweeps: along the columns of every row, then along the rows.
140type Pool is Data:
141  Pool{height: U32,width: U32,kernel: U32,stride: U32,padding: U32,out_height: U32,out_width: U32}
142
143def window_read(~Element: Data,+pad: Element,valid: Bool,input: Array<Element>,index: U32) -> Array<Element> & Element:
144  match valid:
145    case False{}: (input,pad)
146    case True{}: Array.get(Element,input,index)
147
148# Window reduction along one axis. step is the number of elements after the
149# axis, the same in the input and the output. Output position p keeps its
150# step offset p % step, lies at coordinate (p / step) % out_extent of the
151# axis, and folds kernel samples in increasing coordinate from pad; a sample
152# outside the axis reads pad. origin is the coordinate at which the window of
153# output coordinate zero starts before padding: zero for a whole axis, the
154# stored row for the rows of a taller image.
155type Sweep is Data:
156  Sweep{extent: U32,step: U32,kernel: U32,stride: U32,padding: U32,out_extent: U32,origin: U32}
157
158type SweepContext<-Element: Data> is Data:
159  SweepContext{sweep: Sweep,pad: Element}
160
161type Line<-Element: Data> is Data:
162  Line{base: U32,first: U32,extent: U32,step: U32,padding: U32,pad: Element}
163
164def read_line(~Element: Data,input: Array<Element>,line: Line<Element>,position: U32) -> Array<Element> & Element:
165  Line{+base,+first,+extent,+step,+padding,+pad} = line
166  +coordinate = (first + position - padding : U32)
167  window_read(~Element,pad,U32.is_lt(coordinate,extent),input,(base + coordinate * step : U32))
168
169def line_at(~Element: Data,+sweep: Sweep,+pad: Element,+position: U32) -> Line<Element>:
170  Sweep{+extent,+step,kernel,+stride,+padding,+out_extent,+origin} = sweep
171  +along = (position / step : U32)
172  Line{(along / out_extent * extent * step + position % step : U32),(along % out_extent * stride + origin : U32),extent,step,padding,pad}
173
174def read_swept(~Element: Data,~operation: Element -> Element -> Element,input: Array<Element>,context: SweepContext<Element>,position: U32) ->
175  Array<Element> & Element:
176  SweepContext{+sweep,+pad} = context
177  Sweep{extent,step,+kernel,stride,padding,out_extent,origin} = sweep
178  Traversal.fold(~Array<Element>,~Element,~Element,~Line<Element>,~read_line(~Element),~operation,U32.to_nat(kernel),0,input,pad,
179    line_at(~Element,sweep,pad,position))
180
181def swept_index(~Element: Data,context: SweepContext<Element>,position: U32) -> U32:
182  position
183
184def sweep_run(~Element: Data,~operation: Element -> Element -> Element,count: Nat,context: SweepContext<Element>,
185  input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:
186  Traversal.finish(~Array<Element>,~Element,~SweepContext<Element>,
187    Traversal.run(~Array<Element>,~Element,~SweepContext<Element>,~read_swept(~Element,~operation),~swept_index(~Element),count,
188      Traversal.State{input,output,0,context}))
189
190# The sweep along the columns of a pool: every stored row is one line.
191def pool_columns(pool: Pool) -> Sweep:
192  Pool{height,width,kernel,stride,padding,out_height,out_width} = pool
193  Sweep{width,1,kernel,stride,padding,out_width,0}
194
195# The sweep along the rows of the column-swept planes. The stored rows may be
196# a window of a taller image; see Sweep for origin.
197def pool_rows(pool: Pool,origin: U32) -> Sweep:
198  Pool{height,width,kernel,stride,padding,out_height,out_width} = pool
199  Sweep{height,out_width,kernel,stride,padding,out_height,origin}