tensor_cursor.bend
stelliferous@0.0.3.0 · 129 lines · 0 definitions
1# Least-significant axis first. Singleton axes are omitted. The two innermost
2# axes are flat levels, so an element step changes registers only; outer axes
3# stay in a list that is rebuilt only when every level wraps. Every step is a
4# direct template or a tail call, so native code keeps the advance in one loop.
5import Base
6import stelliferous@0.0.3.0/tensor_view.bend as Views
7import stelliferous@0.0.3.0/traversal_array.bend as Traversal
8
9type Axis is Data:
10 Axis{extent: Nat,remaining: Nat,stride: U32,rewind: U32}
11
12def include_axis(active: Bool,+extent: U32,+stride: U32,axes: List<&2,Axis>) -> List<&2,Axis>:
13 match active:
14 case False{}: axes
15 case True{}:
16 +count = U32.to_nat(extent)
17 Axis{count,count,stride,((extent - 1) * stride : U32)} <> axes
18
19def axes(source: List<&2,Views.Axis>,reversed: List<&2,Axis>) -> List<&2,Axis>:
20 match source:
21 case []: reversed
22 case Views.Axis{+extent,stride} <> tail:
23 axes(tail,include_axis(U32.is_lt(1,extent),extent,stride,reversed))
24
25# A level with no remaining count never moves; it stands for an absent axis.
26type Level<-Rest: Data> is Data:
27 Level{axis: Axis,rest: Rest}
28
29def absent() -> Axis:
30 Axis{0n,0n,0,0}
31
32# Outer axes: reset axes are kept reversed and restored after the carry stops.
33def restore(reset: List<&2,Axis>,axes: List<&2,Axis>) -> List<&2,Axis>:
34 match reset:
35 case []: axes
36 case axis <> tail: restore(tail,axis <> axes)
37
38def carry(axes: List<&2,Axis>,address: U32,reset: List<&2,Axis>) -> List<&2,Axis> & U32:
39 match axes:
40 case []: (restore(reset,[]),address)
41 case Axis{extent,0n,stride,rewind} <> tail: (restore(reset,Axis{extent,0n,stride,rewind} <> tail),address)
42 case Axis{+extent,1n,stride,+rewind} <> tail: carry(tail,(address - rewind : U32),Axis{extent,extent,stride,rewind} <> reset)
43 case Axis{extent,1n+1n+remaining,+stride,rewind} <> tail:
44 (restore(reset,Axis{extent,1n+remaining,stride,rewind} <> tail),(address + stride : U32))
45
46def advance_outer(axes: List<&2,Axis>,address: U32) -> List<&2,Axis> & U32:
47 carry(axes,address,[])
48
49def reset(~Rest: Data,+extent: Nat,stride: U32,rewind: U32,result: Rest & U32) -> Level<Rest> & U32:
50 (rest,address) = result
51 (Level{Axis{extent,extent,stride,rewind},rest},address)
52
53def step(~Rest: Data,~next: Rest -> U32 -> Rest & U32,level: Level<Rest>,address: U32) -> Level<Rest> & U32:
54 Level{axis,rest} = level
55 match axis:
56 case Axis{extent,0n,stride,rewind}: (Level{Axis{extent,0n,stride,rewind},rest},address)
57 case Axis{+extent,1n,stride,+rewind}: reset(~Rest,extent,stride,rewind,next(rest,(address - rewind : U32)))
58 case Axis{extent,1n+1n+remaining,+stride,rewind}: (Level{Axis{extent,1n+remaining,stride,rewind},rest},(address + stride : U32))
59
60def Depth0() -> Data:
61 List<&2,Axis>
62
63def Depth1() -> Data:
64 Level<Depth0()>
65
66def Levels() -> Data:
67 Level<Depth1()>
68
69def advance_levels(axes: Levels(),address: U32) -> Levels() & U32:
70 step(~Depth1(),~step(~Depth0(),~advance_outer),axes,address)
71
72type Cursor is Data:
73 Cursor{axes: Levels(),address: U32}
74
75def moved(result: Levels() & U32) -> Cursor:
76 (axes,address) = result
77 Cursor{axes,address}
78
79def advance(axes: Levels(),address: U32) -> Cursor:
80 moved(advance_levels(axes,address))
81
82def split_level(~Rest: Data,~split_rest: List<&2,Axis> -> Rest,axes: List<&2,Axis>) -> Level<Rest>:
83 match axes:
84 case []: Level{absent(),split_rest([])}
85 case axis <> tail: Level{axis,split_rest(tail)}
86
87def outer(axes: List<&2,Axis>) -> List<&2,Axis>:
88 axes
89
90def split(axes: List<&2,Axis>) -> Levels():
91 split_level(~Depth1(),~split_level(~Depth0(),~outer),axes)
92
93def prepare(view: Views.View) -> Cursor:
94 Views.View{source,offset} = view
95 Cursor{split(axes(source,[])),offset}
96
97def read_value(~Element: Data,axes: Levels(),address: U32,result: Array<Element> & Element) -> Array<Element> & Traversal.Visited<Cursor,Element>:
98 (storage,value) = result
99 (storage,Traversal.Visited{advance(axes,address),value})
100
101def read(~Element: Data,storage: Array<Element>,cursor: Cursor) -> Array<Element> & Traversal.Visited<Cursor,Element>:
102 Cursor{axes,+address} = cursor
103 read_value(~Element,axes,address,Array.get(Element,storage,address))
104
105# One run of the innermost axis reads address, address + stride, ... with no
106# countdown. Its two words stay in registers in whatever program the cursor is
107# compiled into; the full cursor is advanced once, at the end of the run.
108type Stride is Data:
109 Stride{address: U32,stride: U32}
110
111def stride_value(~Element: Data,address: U32,+stride: U32,result: Array<Element> & Element) -> Array<Element> & Traversal.Visited<Stride,Element>:
112 (storage,value) = result
113 (storage,Traversal.Visited{Stride{(address + stride : U32),stride},value})
114
115def read_stride(~Element: Data,storage: Array<Element>,cursor: Stride) -> Array<Element> & Traversal.Visited<Stride,Element>:
116 Stride{+address,+stride} = cursor
117 stride_value(~Element,address,stride,Array.get(Element,storage,address))
118
119# The address a scatter stores at, with the cursor of the next element.
120def target(cursor: Cursor) -> Traversal.Target<Cursor>:
121 Cursor{axes,+address} = cursor
122 Traversal.Target{advance(axes,address),address}
123
124def target_stride(cursor: Stride) -> Traversal.Target<Stride>:
125 Stride{+address,+stride} = cursor
126 Traversal.Target{Stride{(address + stride : U32),stride},address}
127
128type CursorPair is Data:
129 CursorPair{left: Cursor,right: Cursor}