traversal_lanes.bend
stelliferous@0.0.3.0 · 121 lines · 0 definitions
1# The FP32 cells of Cells.map: eight consecutive values as one Tile8.
2import Base
3import stelliferous@0.0.3.0/kernel_tile.bend as Tile
4import stelliferous@0.0.3.0/traversal_cells.bend as Cells
5import stelliferous@0.0.3.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)
102
103# An arbitrary gather uses the common rounds and the operation's scalar tail.
104def map_with(~Context: Data,~operation: Context -> F32 -> F32,~collect: Array<F32> -> U32 -> Context -> Array<F32> & Tile.Tile8,count: Nat,position: U32,context: Context,values: Array<F32>) -> Array<F32>:
105 rounds_with(~Array<F32>,~Context,~collect,~scatter,~finish(~Context,~operation),count,Cells.State{values,scratch(),position,0,context})
106
107def map_with_when(~Context: Data,~operation: Context -> F32 -> F32,~collect: Array<F32> -> U32 -> Context -> Array<F32> & Tile.Tile8,covered: Bool,count: Nat,position: U32,context: Context,values: Array<F32>) -> Array<F32>:
108 match covered:
109 case True{}: map_with(~Context,~operation,~collect,count,position,context,values)
110 case False{}: Traversal.map_in_place(~F32,~Context,~operation,count,position,context,values)
111
112def identity(~Context: Data,context: Context,value: F32) -> F32:
113 value
114
115def transformed(~Context: Data,~transform: Context -> Tile.Tile8 -> Tile.Tile8,context: Context,read: Array<F32> & Tile.Tile8) -> Array<F32> & Tile.Tile8:
116 (values,tile) = read
117 (values,transform(context,tile))
118
119def gather_transformed(~Context: Data,~transform: Context -> Tile.Tile8 -> Tile.Tile8,values: Array<F32>,position: U32,context: Context) -> Array<F32> & Tile.Tile8:
120 +context = context
121 transformed(~Context,~transform,context,gather(~Context,~identity(~Context),values,position,context))