1	# Internal tensor read layouts and unchecked traversal kernels.
2	import Base
3	import ./tensor_view.bend as Views
4	import ./traversal_array.bend as Traversal
5	import ./traversal_loop.bend as Loop
6	import ./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.
15	def 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	
20	def 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	
29	def 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.
37	def 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	
48	def 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	
60	def 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	
65	def 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).
71	def 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].
77	def 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	
82	def 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	
85	def 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	
93	def zero_contiguous(+view: Views.View) -> Bool:
94	  Views.View{axes,offset} = view
95	  U32.is_zero(offset) && Views.is_contiguous(view)
96	
97	def 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	
106	def 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.
114	type Runs is Data:
115	  Runs{cursor: Cursor.Cursor,extent: U32,stride: U32}
116	
117	type Lines<-Accumulator: Data> is Data:
118	  Lines{cursor: Cursor.Cursor,runs: Nat,extent: U32,stride: U32,initial: Accumulator}
119	
120	def 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.
126	def 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	
133	def 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	
138	def 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.
146	def 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.
154	type Pool is Data:
155	  Pool{height: U32,width: U32,kernel: U32,stride: U32,padding: U32,out_height: U32,out_width: U32}
156	
157	def 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.
169	type Sweep is Data:
170	  Sweep{extent: U32,step: U32,kernel: U32,stride: U32,padding: U32,out_extent: U32,origin: U32}
171	
172	type SweepContext<-Element: Data> is Data:
173	  SweepContext{sweep: Sweep,pad: Element}
174	
175	type Line<-Element: Data> is Data:
176	  Line{base: U32,first: U32,extent: U32,step: U32,padding: U32,pad: Element}
177	
178	def 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	
183	def 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	
188	def 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	
195	def swept_index(~Element: Data,context: SweepContext<Element>,position: U32) -> U32:
196	  position
197	
198	def 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.
205	def 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.
211	def 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}