kernel_tile.bend
stelliferous@0.0.2.0 · 96 lines · 19 definitions
1
2import Base
3import stelliferous@0.0.2.0/kernel_block.bend as Block
4import stelliferous@0.0.2.0/math_sequence.bend as Sequence
5import stelliferous@0.0.2.0/math_fp32.bend as Fp32
6
7type Tile8 is Data:
8 Tile{a0: F32, a1: F32, a2: F32, a3: F32, a4: F32, a5: F32, a6: F32, a7: F32}
9
10def filled(+value: F32) -> Tile8:
11 Tile{value,value,value,value,value,value,value,value}
12
13def lane0(tile: Tile8) -> F32:
14 Tile{a,b,c,d,e,f,g,h} = tile
15 a
16
17def lane1(tile: Tile8) -> F32:
18 Tile{a,b,c,d,e,f,g,h} = tile
19 b
20
21def lane2(tile: Tile8) -> F32:
22 Tile{a,b,c,d,e,f,g,h} = tile
23 c
24
25def lane3(tile: Tile8) -> F32:
26 Tile{a,b,c,d,e,f,g,h} = tile
27 d
28
29def lane4(tile: Tile8) -> F32:
30 Tile{a,b,c,d,e,f,g,h} = tile
31 e
32
33def lane5(tile: Tile8) -> F32:
34 Tile{a,b,c,d,e,f,g,h} = tile
35 f
36
37def lane6(tile: Tile8) -> F32:
38 Tile{a,b,c,d,e,f,g,h} = tile
39 g
40
41def lane7(tile: Tile8) -> F32:
42 Tile{a,b,c,d,e,f,g,h} = tile
43 h
44
45
46def update(accumulator: Tile8,input: Tile8,weight: F32) -> Tile8:
47 Tile{a,b,c,d,e,f,g,h} = accumulator
48 Tile{x0,x1,x2,x3,x4,x5,x6,x7} = input
49 +weight = weight
50 Tile{Fp32.multiply_add(a,x0,weight),Fp32.multiply_add(b,x1,weight),Fp32.multiply_add(c,x2,weight),Fp32.multiply_add(d,x3,weight),
51 Fp32.multiply_add(e,x4,weight),Fp32.multiply_add(f,x5,weight),Fp32.multiply_add(g,x6,weight),Fp32.multiply_add(h,x7,weight)}
52
53def from_sequence(vector: Sequence.Values(F32,8n)) -> Tile8:
54 Sequence.Elements{a0,Sequence.Elements{a1,Sequence.Elements{a2,Sequence.Elements{a3,Sequence.Elements{a4,Sequence.Elements{a5,Sequence.Elements{a6,
55 Sequence.Elements{a7,Unit{}}}}}}}}} = vector
56 Tile{a0,a1,a2,a3,a4,a5,a6,a7}
57
58type HeadContext<-Context: Data> is Data:
59 HeadContext{head: F32,context: Context}
60
61def read_with_head(
62 ~State: Type,~Context: Data,~read: State -> Context -> Nat -> State & F32,
63 state: State,combined: HeadContext<Context>,index: Nat
64) -> State & F32:
65 HeadContext{head,context} = combined
66 match index:
67 case 0n: (state,head)
68 case 1n+rest: read(state,context,1n+rest)
69
70def pairs_value(vector: Block.Twin<Block.Twin<Block.Twin<F32>>>) -> Tile8:
71 Block.Twin{Block.Twin{Block.Twin{a0,a1},Block.Twin{a2,a3}},Block.Twin{Block.Twin{a4,a5},
72 Block.Twin{a6,a7}}} = vector
73 Tile{a0,a1,a2,a3,a4,a5,a6,a7}
74
75def finish_pairs(~State: Type,result: State & Block.Twin<Block.Twin<Block.Twin<F32>>>) -> State & Tile8:
76 (state,vector) = result
77 (state,pairs_value(vector))
78
79def read_tail(~State: Type,~Context: Data,~read: State -> Context -> Nat -> State & F32,first: State & F32,context: Context) -> State & Tile8:
80 (state,head) = first
81 finish_pairs(~State,
82 Block.read_pair(~State,~HeadContext<Context>,~Block.Twin<Block.Twin<F32>>,
83 ~Block.read_pair(~State,~HeadContext<Context>,~Block.Twin<F32>,
84 ~Block.read_pair(~State,~HeadContext<Context>,~F32,
85 ~read_with_head(~State,~Context,~read),~1n),~2n),~4n,
86 state,HeadContext{head,context},0n))
87
88
89
90def view(tile: Tile8) -> Sequence.Values(F32,8n):
91 Tile{a0,a1,a2,a3,a4,a5,a6,a7} = tile
92 Sequence.Elements{a0,Sequence.Elements{a1,Sequence.Elements{a2,Sequence.Elements{a3,Sequence.Elements{a4,Sequence.Elements{a5,Sequence.Elements{a6,
93 Sequence.Elements{a7,Unit{}}}}}}}}}
94
95def element_at(index: Sequence.Index(8n),tile: Tile8) -> F32:
96 Sequence.get(F32,8n,index,view(tile))