kernel_weights.bend
stelliferous@0.0.3.0 · 129 lines · 0 definitions
1import Base
2import stelliferous@0.0.3.0/storage_buffer.bend as Storage
3import stelliferous@0.0.3.0/storage_allocation.bend as Allocation
4import stelliferous@0.0.3.0/traversal_array.bend as Traversal
5import stelliferous@0.0.3.0/kernel_tile.bend as Tile
6import stelliferous@0.0.3.0/tensor.bend as Tensors
7
8# Each ordered K coordinate owns eight adjacent weights, one per output row of
9# a block; rows past the last output read zero. One native tile read replaces
10# eight separately masked scalar reads.
11type MappedLayout<-Context: Data> is Data:
12 MappedLayout{rows: U32,inner: U32,row_stride: U32,context: Context}
13
14def identity(context: Unit,index: U32) -> U32:
15 index
16
17type Bases is Data:
18 Bases{a: U32,b: U32,c: U32,d: U32,e: U32,f: U32,g: U32,h: U32}
19
20def linear_bases(~Context: Data,context: Context,+row_stride: U32,+row: U32) -> Bases:
21 Bases{(row * row_stride : U32),((row + 1) * row_stride : U32),((row + 2) * row_stride : U32),((row + 3) * row_stride : U32),
22 ((row + 4) * row_stride : U32),((row + 5) * row_stride : U32),((row + 6) * row_stride : U32),((row + 7) * row_stride : U32)}
23
24def base0(bases: Bases) -> U32:
25 Bases{a,b,c,d,e,f,g,h} = bases
26 a
27
28def base1(bases: Bases) -> U32:
29 Bases{a,b,c,d,e,f,g,h} = bases
30 b
31
32def base2(bases: Bases) -> U32:
33 Bases{a,b,c,d,e,f,g,h} = bases
34 c
35
36def base3(bases: Bases) -> U32:
37 Bases{a,b,c,d,e,f,g,h} = bases
38 d
39
40def base4(bases: Bases) -> U32:
41 Bases{a,b,c,d,e,f,g,h} = bases
42 e
43
44def base5(bases: Bases) -> U32:
45 Bases{a,b,c,d,e,f,g,h} = bases
46 f
47
48def base6(bases: Bases) -> U32:
49 Bases{a,b,c,d,e,f,g,h} = bases
50 g
51
52def base7(bases: Bases) -> U32:
53 Bases{a,b,c,d,e,f,g,h} = bases
54 h
55
56def eighth(a: F32,b: F32,c: F32,d: F32,e: F32,f: F32,g: F32,read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
57 (input,h) = read
58 (input,Tile.Tile{a,b,c,d,e,f,g,h})
59
60def seventh_rows(+rows: U32,+row: U32,+index: U32,+bases: Bases,a: F32,b: F32,c: F32,d: F32,e: F32,f: F32,read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
61 (input,g) = read
62 eighth(a,b,c,d,e,f,g,Storage.read_when(~F32,0.0,U32.is_lt((row + 7 : U32),rows),input,(base7(bases) + index : U32)))
63
64def sixth_rows(+rows: U32,+row: U32,+index: U32,+bases: Bases,a: F32,b: F32,c: F32,d: F32,e: F32,read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
65 (input,f) = read
66 seventh_rows(rows,row,index,bases,a,b,c,d,e,f,Storage.read_when(~F32,0.0,U32.is_lt((row + 6 : U32),rows),input,(base6(bases) + index : U32)))
67
68def fifth_rows(+rows: U32,+row: U32,+index: U32,+bases: Bases,a: F32,b: F32,c: F32,d: F32,read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
69 (input,e) = read
70 sixth_rows(rows,row,index,bases,a,b,c,d,e,Storage.read_when(~F32,0.0,U32.is_lt((row + 5 : U32),rows),input,(base5(bases) + index : U32)))
71
72def fourth_rows(+rows: U32,+row: U32,+index: U32,+bases: Bases,a: F32,b: F32,c: F32,read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
73 (input,d) = read
74 fifth_rows(rows,row,index,bases,a,b,c,d,Storage.read_when(~F32,0.0,U32.is_lt((row + 4 : U32),rows),input,(base4(bases) + index : U32)))
75
76def third_rows(+rows: U32,+row: U32,+index: U32,+bases: Bases,a: F32,b: F32,read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
77 (input,c) = read
78 fourth_rows(rows,row,index,bases,a,b,c,Storage.read_when(~F32,0.0,U32.is_lt((row + 3 : U32),rows),input,(base3(bases) + index : U32)))
79
80def second_rows(+rows: U32,+row: U32,+index: U32,+bases: Bases,a: F32,read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
81 (input,b) = read
82 third_rows(rows,row,index,bases,a,b,Storage.read_when(~F32,0.0,U32.is_lt((row + 2 : U32),rows),input,(base2(bases) + index : U32)))
83
84def first_rows(+rows: U32,+row: U32,+index: U32,+bases: Bases,read: Array<F32> & F32) -> Array<F32> & Tile.Tile8:
85 (input,a) = read
86 second_rows(rows,row,index,bases,a,Storage.read_when(~F32,0.0,U32.is_lt((row + 1 : U32),rows),input,(base1(bases) + index : U32)))
87
88def read_rows(input: Array<F32>,+rows: U32,+row: U32,+index: U32,+bases: Bases) -> Array<F32> & Tile.Tile8:
89 first_rows(rows,row,index,bases,Storage.read_when(~F32,0.0,U32.is_lt(row,rows),input,(base0(bases) + index : U32)))
90
91type MappedRows<-Context: Data> is Data:
92 MappedRows{rows: U32,row: U32,destination: U32,context: Context,bases: Bases}
93
94def read_mapped_rows(~Context: Data,~offset: Context -> U32 -> U32,input: Array<F32>,context: MappedRows<Context>,index: U32) -> Array<F32> & Tile.Tile8:
95 MappedRows{rows,row,destination,context,bases} = context
96 read_rows(input,rows,row,offset(context,index),bases)
97
98def destination_rows(~Context: Data,context: MappedRows<Context>,index: U32) -> U32:
99 MappedRows{rows,row,destination,context,bases} = context
100 (destination + index : U32)
101
102def copy_mapped_with(~Context: Data,~offset: Context -> U32 -> U32,~row_bases: Context -> @+row_stride: U32 -> @+row: U32 -> Bases,block: U32,layout: MappedLayout<Context>,owners: Array<F32> & Array<Tile.Tile8>) -> Array<F32> & Array<Tile.Tile8>:
103 MappedLayout{rows,+inner,row_stride,+context} = layout
104 (input,output) = owners
105 +block = block
106 +row = (block * 8 : U32)
107 +bases = row_bases(context,row_stride,row)
108 Traversal.finish(~Array<F32>,~Tile.Tile8,~MappedRows<Context>,
109 Traversal.run(~Array<F32>,~Tile.Tile8,~MappedRows<Context>,~read_mapped_rows(~Context,~offset),~destination_rows(~Context),U32.to_nat(inner),
110 Traversal.State{input,output,0,MappedRows{rows,row,(block * inner : U32),context,bases}}))
111
112def source(pair: Array<F32> & Array<Tile.Tile8>) -> Array<F32>:
113 (input,output) = pair
114 input
115
116def reordered(pair: Array<F32> & Array<Tile.Tile8>) -> Array<Tile.Tile8>:
117 (input,output) = pair
118 output
119
120def run_mapped_with(~Context: Data,~offset: Context -> U32 -> U32,~row_bases: Context -> @+row_stride: U32 -> @+row: U32 -> Bases,+rows: U32,+inner: U32,row_stride: U32,context: Context,input: Array<F32>) -> Array<F32> & Array<Tile.Tile8>:
121 +blocks = ((rows + 7) / 8 : U32)
122 Traversal.indexed_fold(~(Array<F32> & Array<Tile.Tile8>),~MappedLayout<Context>,~copy_mapped_with(~Context,~offset,~row_bases),U32.to_nat(blocks),0,
123 MappedLayout{rows,inner,row_stride,context},(input,Allocation.array(~Tile.Tile8,Tensors.depth_for((blocks * inner : U32)),Tile.filled(0.0))))
124
125def run_mapped(~Context: Data,~offset: Context -> U32 -> U32,+rows: U32,+inner: U32,row_stride: U32,context: Context,input: Array<F32>) -> Array<F32> & Array<Tile.Tile8>:
126 run_mapped_with(~Context,~offset,~linear_bases(~Context),rows,inner,row_stride,context,input)
127
128def run(+rows: U32,+inner: U32,input: Array<F32>) -> Array<Tile.Tile8>:
129 reordered(run_mapped(~Unit,~identity,rows,inner,inner,Unit{},input))