kernel_weights.bend
stelliferous@0.0.2.0 · 73 lines · 15 definitions
1import Base
2import stelliferous@0.0.2.0/storage_buffer.bend as Storage
3import stelliferous@0.0.2.0/storage_allocation.bend as Allocation
4import stelliferous@0.0.2.0/traversal_array.bend as Traversal
5import stelliferous@0.0.2.0/kernel_tile.bend as Tile
6import stelliferous@0.0.2.0/tensor.bend as Tensors
7
8
9
10
11type Layout is Data:
12 Layout{rows: U32,inner: U32}
13
14type Context is Data:
15 Context{rows: U32,inner: U32,row: U32,destination: U32}
16
17def eighth(a: F32,b: F32,c: F32,d: F32,e: F32,f: F32,g: F32,read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
18 (input,h) = read
19 (input,Tile.Tile{a,b,c,d,e,f,g,h})
20
21def seventh(+rows: U32,+inner: U32,+row: U32,+index: U32,a: F32,b: F32,c: F32,d: F32,e: F32,f: F32,read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
22 (input,g) = read
23 eighth(a,b,c,d,e,f,g,Storage.read_when(~F32,0.0,U32.is_lt((row + 7 : U32),rows),input,((row + 7) * inner + index : U32)))
24
25def sixth(+rows: U32,+inner: U32,+row: U32,+index: U32,a: F32,b: F32,c: F32,d: F32,e: F32,read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
26 (input,f) = read
27 seventh(rows,inner,row,index,a,b,c,d,e,f,Storage.read_when(~F32,0.0,U32.is_lt((row + 6 : U32),rows),input,((row + 6) * inner + index : U32)))
28
29def fifth(+rows: U32,+inner: U32,+row: U32,+index: U32,a: F32,b: F32,c: F32,d: F32,read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
30 (input,e) = read
31 sixth(rows,inner,row,index,a,b,c,d,e,Storage.read_when(~F32,0.0,U32.is_lt((row + 5 : U32),rows),input,((row + 5) * inner + index : U32)))
32
33def fourth(+rows: U32,+inner: U32,+row: U32,+index: U32,a: F32,b: F32,c: F32,read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
34 (input,d) = read
35 fifth(rows,inner,row,index,a,b,c,d,Storage.read_when(~F32,0.0,U32.is_lt((row + 4 : U32),rows),input,((row + 4) * inner + index : U32)))
36
37def third(+rows: U32,+inner: U32,+row: U32,+index: U32,a: F32,b: F32,read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
38 (input,c) = read
39 fourth(rows,inner,row,index,a,b,c,Storage.read_when(~F32,0.0,U32.is_lt((row + 3 : U32),rows),input,((row + 3) * inner + index : U32)))
40
41def second(+rows: U32,+inner: U32,+row: U32,+index: U32,a: F32,read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
42 (input,b) = read
43 third(rows,inner,row,index,a,b,Storage.read_when(~F32,0.0,U32.is_lt((row + 2 : U32),rows),input,((row + 2) * inner + index : U32)))
44
45def first(+rows: U32,+inner: U32,+row: U32,+index: U32,read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
46 (input,a) = read
47 second(rows,inner,row,index,a,Storage.read_when(~F32,0.0,U32.is_lt((row + 1 : U32),rows),input,((row + 1) * inner + index : U32)))
48
49def read(input: Array<F32>,context: Context,index: U32) -> Array<F32> & Tile.Tile8:
50 Context{+rows,+inner,+row,destination} = context
51 +index = index
52 first(rows,inner,row,index,Storage.read_when(~F32,0.0,U32.is_lt(row,rows),input,(row * inner + index : U32)))
53
54def destination(context: Context,index: U32) -> U32:
55 Context{rows,inner,row,destination} = context
56 (destination + index : U32)
57
58def copy_block(block: U32,layout: Layout,owners: Array<F32> & Array<Tile.Tile8>) -> Array<F32> & Array<Tile.Tile8>:
59 Layout{rows,+inner} = layout
60 (input,output) = owners
61 +block = block
62 Traversal.finish(~Array<F32>,~Tile.Tile8,~Context,
63 Traversal.run(~Array<F32>,~Tile.Tile8,~Context,~read,~destination,U32.to_nat(inner),
64 Traversal.State{input,output,0,Context{rows,inner,(block * 8 : U32),(block * inner : U32)}}))
65
66def reordered(pair: Array<F32> & Array<Tile.Tile8>) -> Array<Tile.Tile8>:
67 (input,output) = pair
68 output
69
70def run(+rows: U32,+inner: U32,input: Array<F32>) -> Array<Tile.Tile8>:
71 +blocks = ((rows + 7) / 8 : U32)
72 reordered(Traversal.indexed_fold(~(Array<F32> & Array<Tile.Tile8>),~Layout,~copy_block,U32.to_nat(blocks),0,
73 Layout{rows,inner},(input,Allocation.array(~Tile.Tile8,Tensors.depth_for((blocks * inner : U32)),Tile.filled(0.0)))))