traversal_lanes.bend

stelliferous@0.0.2.0 · 101 lines · 17 definitions

1# The FP32 cells of Cells.map: eight consecutive values as one Tile8.
2import Base
3import stelliferous@0.0.2.0/kernel_tile.bend as Tile
4import stelliferous@0.0.2.0/traversal_cells.bend as Cells
5import stelliferous@0.0.2.0/traversal_array.bend as Traversal
6
7def gathered(~Context: Data,~map: Context -> F32 -> F32,+context: Context,+a: F32,+b: F32,+c: F32,+d: F32,+e: F32,+f: F32,+g: F32,
8  read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
9  (values,h) = read
10  (values,Tile.Tile{map(context,a),map(context,b),map(context,c),map(context,d),map(context,e),map(context,f),map(context,g),map(context,h)})
11
12def gather7(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,+a: F32,+b: F32,+c: F32,+d: F32,+e: F32,+f: F32,
13  read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
14  (values,g) = read
15  gathered(~Context,~map,context,a,b,c,d,e,f,g,Array.get(F32,values,(position + 1 : U32)))
16
17def gather6(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,+a: F32,+b: F32,+c: F32,+d: F32,+e: F32,
18  read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
19  (values,f) = read
20  gather7(~Context,~map,(position + 1 : U32),context,a,b,c,d,e,f,Array.get(F32,values,(position + 1 : U32)))
21
22def gather5(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,+a: F32,+b: F32,+c: F32,+d: F32,
23  read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
24  (values,e) = read
25  gather6(~Context,~map,(position + 1 : U32),context,a,b,c,d,e,Array.get(F32,values,(position + 1 : U32)))
26
27def gather4(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,+a: F32,+b: F32,+c: F32,
28  read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
29  (values,d) = read
30  gather5(~Context,~map,(position + 1 : U32),context,a,b,c,d,Array.get(F32,values,(position + 1 : U32)))
31
32def gather3(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,+a: F32,+b: F32,
33  read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
34  (values,c) = read
35  gather4(~Context,~map,(position + 1 : U32),context,a,b,c,Array.get(F32,values,(position + 1 : U32)))
36
37def gather2(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,+a: F32,
38  read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
39  (values,b) = read
40  gather3(~Context,~map,(position + 1 : U32),context,a,b,Array.get(F32,values,(position + 1 : U32)))
41
42def gather1(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,
43  read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
44  (values,a) = read
45  gather2(~Context,~map,(position + 1 : U32),context,a,Array.get(F32,values,(position + 1 : U32)))
46
47# The mapped values at position, position + 1, ... as one tile; the array is only read.
48def gather(~Context: Data,~map: Context -> F32 -> F32,values: Array<F32>,position: U32,context: Context) -> Array<F32> & Tile.Tile8:
49  +position = position
50  gather1(~Context,~map,position,context,Array.get(F32,values,position))
51
52# The lanes of a tile to position, position + 1, ...
53def scatter(values: Array<F32>,position: U32,tile: Tile.Tile8) -> Array<F32>:
54  +tile = tile
55  +position = position
56  +p1 = (position + 1 : U32)
57  +p2 = (p1 + 1 : U32)
58  +p3 = (p2 + 1 : U32)
59  +p4 = (p3 + 1 : U32)
60  +p5 = (p4 + 1 : U32)
61  +p6 = (p5 + 1 : U32)
62  Array.set(F32,Array.set(F32,Array.set(F32,Array.set(F32,Array.set(F32,Array.set(F32,Array.set(F32,Array.set(F32,values,position,Tile.lane0(tile)),p1,Tile.lane1(tile)),
63    p2,Tile.lane2(tile)),p3,Tile.lane3(tile)),p4,Tile.lane4(tile)),p5,Tile.lane5(tile)),p6,Tile.lane6(tile)),(p6 + 1 : U32),Tile.lane7(tile))
64
65def scratch() -> Array<Tile.Tile8>:
66  Array.new(Tile.Tile8,6n,Tile.filled(0.0))
67
68# The values the rounds leave: fewer than eight, mapped one by one.
69def finish(~Context: Data,~map: Context -> F32 -> F32,rest: Nat,state: Cells.State<Array<F32>,Tile.Tile8,Context>) -> Array<F32>:
70  Cells.State{values,scratch,position,slot,context} = state
71  Traversal.map_in_place(~F32,~Context,~map,rest,position,context,values)
72
73# One eight-value loop for any owned input bundle; gather and scatter decide
74# which owners are read or updated. The unary path is its Array specialization.
75def rounds_with(~Owner: Type,~Context: Data,~gather: Owner -> U32 -> Context -> Owner & Tile.Tile8,
76  ~scatter: Owner -> U32 -> Tile.Tile8 -> Owner,~finish: Nat -> Cells.State<Owner,Tile.Tile8,Context> -> Owner,
77  count: Nat,state: Cells.State<Owner,Tile.Tile8,Context>) -> Owner:
78  match count:
79    case 1n+1n+1n+1n+1n+1n+1n+1n+rest:
80      rounds_with(~Owner,~Context,~gather,~scatter,~finish,rest,Cells.step(~Owner,~Tile.Tile8,~Context,~gather,~scatter,~8,state))
81    case few: finish(few,state)
82
83def rounds(~Context: Data,~map: Context -> F32 -> F32,count: Nat,state: Cells.State<Array<F32>,Tile.Tile8,Context>) -> Array<F32>:
84  rounds_with(~Array<F32>,~Context,~gather(~Context,~map),~scatter,~finish(~Context,~map),count,state)
85
86# map over the count values from position, in order: Traversal.map_in_place
87# with the values of each round computed together.
88def map(~Context: Data,~map: Context -> F32 -> F32,count: Nat,position: U32,context: Context,values: Array<F32>) -> Array<F32>:
89  rounds(~Context,~map,count,Cells.State{values,scratch(),position,0,context})
90
91# An array of 2^depth values holds the count positions from position, within
92# the depths the refinement proof covers.
93def covered(+depth: Nat,+position: U32,+count: Nat) -> Bool:
94  Nat.is_le(depth,30n) && Nat.is_le(Nat.add(U32.to_nat(position),count),U32.to_nat(U32.shln(1,depth)))
95
96# map over count values from position when the array is known to hold them,
97# and the scalar map otherwise: the same array either way.
98def map_when(~Context: Data,~operation: Context -> F32 -> F32,covered: Bool,count: Nat,position: U32,context: Context,values: Array<F32>) -> Array<F32>:
99  match covered:
100    case True{}: map(~Context,~operation,count,position,context,values)
101    case False{}: Traversal.map_in_place(~F32,~Context,~operation,count,position,context,values)