kernel_gemm.bend
stelliferous@0.0.3.0 · 200 lines · 0 definitions
1import Base
2import stelliferous@0.0.3.0/kernel_rows.bend as Rows
3import stelliferous@0.0.3.0/kernel_tile.bend as Tile
4import stelliferous@0.0.3.0/traversal_loop.bend as Iteration
5import stelliferous@0.0.3.0/traversal_array.bend as Traversal
6import stelliferous@0.0.3.0/traversal_lanes.bend as Lanes
7
8type Buffers is Type:
9 Buffers{input: Array<Tile.Tile8>,weights: Array<Tile.Tile8>,bias: Array<F32>,output: Array<F32>}
10
11# Row-major output, eight contiguous columns; packed weights are [M/8,K,8].
12# A job is one half of a weight block, four rows, at eight positions. Only
13# valid rows are stored; the last weight block is zero-padded to eight rows.
14type Layout is Data:
15 Layout{rows: U32,columns: U32,inner: U32,start: U32,output_columns: U32}
16
17# input_base and weight_base address the job's first input group and weight
18# tile; the loop index advances both.
19type Context is Data:
20 Context{row: U32,position: U32,rows: U32,columns: U32,input_base: U32,weight_base: U32}
21
22type Job is Data:
23 Job{context: Context,upper: Bool,count: Nat}
24
25# One flat loop state. Every owner, the index and the accumulators are fields
26# of this record, and each helper that passes the accumulators on carries few
27# other values: a nested owner record or a longer helper chain made every k
28# step a real call.
29type Loop is Type:
30 Loop{input: Array<Tile.Tile8>,weights: Array<Tile.Tile8>,bias: Array<F32>,output: Array<F32>,index: U32,context: Context,
31 accumulators: Rows.Rows<Tile.Tile8>}
32
33# A job owns four output rows: one half of an eight-row weight block.
34def row_units(rows: U32) -> U32:
35 ((rows + 3) / 4 : U32)
36
37def column_groups(columns: U32) -> U32:
38 (columns / 8 : U32)
39
40def job_count(rows: U32,columns: U32) -> U32:
41 (column_groups(columns) * row_units(rows) : U32)
42
43# Input-major: one input group stays in cache while every row unit passes
44# over it (the BLIS loop order). Units 2b and 2b+1 share weight block b.
45def job(layout: Layout,index: U32) -> Job:
46 Layout{+rows,columns,+inner,+start,+output_columns} = layout
47 +index = index
48 +units = row_units(rows)
49 +unit = (index % units : U32)
50 +group = (index / units : U32)
51 Job{Context{(unit * 4 : U32),(start + group * 8 : U32),rows,output_columns,(group * inner : U32),(unit / 2 * inner : U32)},U32.is_eq((unit % 2 : U32),1),U32.to_nat(inner)}
52
53def update(accumulators: Rows.Rows<Tile.Tile8>,+input: Tile.Tile8,weights: Rows.Rows<F32>) -> Rows.Rows<Tile.Tile8>:
54 Rows.map(~Tile.Tile8,~Tile.Tile8,~F32,~Tile.update,accumulators,input,weights)
55
56def with_input(weights: Array<Tile.Tile8>,bias: Array<F32>,output: Array<F32>,+index: U32,context: Context,
57 accumulators: Rows.Rows<Tile.Tile8>,samples: Rows.Rows<F32>,read: Array<Tile.Tile8> & Tile.Tile8) -> Loop:
58 (input,tile) = read
59 Loop{input,weights,bias,output,(index + 1 : U32),context,update(accumulators,tile,samples)}
60
61def with_weights(~upper: Bool,input: Array<Tile.Tile8>,bias: Array<F32>,output: Array<F32>,+index: U32,context: Context,+input_base: U32,
62 accumulators: Rows.Rows<Tile.Tile8>,read: Array<Tile.Tile8> & Tile.Tile8) -> Loop:
63 (weights,tile) = read
64 with_input(weights,bias,output,index,context,accumulators,Rows.half(upper,tile),Array.get(Tile.Tile8,input,(input_base + index : U32)))
65
66# One ordered k step: read the weight tile, then the input tile, then MAC.
67def step(~upper: Bool,state: Loop) -> Loop:
68 Loop{input,weights,bias,output,+index,context,accumulators} = state
69 Context{+row,+position,+rows,+columns,+input_base,+weight_base} = context
70 with_weights(~upper,input,bias,output,index,Context{row,position,rows,columns,input_base,weight_base},input_base,accumulators,
71 Array.get(Tile.Tile8,weights,(weight_base + index : U32)))
72
73# The last group can contain padding lanes; only real columns are written.
74def store_lane(inside: Bool,output: Array<F32>,offset: U32,value: F32) -> Array<F32>:
75 match inside:
76 case True{}: Array.set(F32,output,offset,value)
77 case False{}: output
78
79def store_tail(+biased: Bool,+count: U32,output: Array<F32>,+offset: U32,+bias: F32,+accumulator: Tile.Tile8) -> Array<F32>:
80 output = store_lane(U32.is_lt(0,count),output,offset,Bool.pick(F32,biased,(Tile.lane0(accumulator) + bias : F32),Tile.lane0(accumulator)))
81 output = store_lane(U32.is_lt(1,count),output,(offset + 1 : U32),Bool.pick(F32,biased,(Tile.lane1(accumulator) + bias : F32),Tile.lane1(accumulator)))
82 output = store_lane(U32.is_lt(2,count),output,(offset + 2 : U32),Bool.pick(F32,biased,(Tile.lane2(accumulator) + bias : F32),Tile.lane2(accumulator)))
83 output = store_lane(U32.is_lt(3,count),output,(offset + 3 : U32),Bool.pick(F32,biased,(Tile.lane3(accumulator) + bias : F32),Tile.lane3(accumulator)))
84 output = store_lane(U32.is_lt(4,count),output,(offset + 4 : U32),Bool.pick(F32,biased,(Tile.lane4(accumulator) + bias : F32),Tile.lane4(accumulator)))
85 output = store_lane(U32.is_lt(5,count),output,(offset + 5 : U32),Bool.pick(F32,biased,(Tile.lane5(accumulator) + bias : F32),Tile.lane5(accumulator)))
86 output = store_lane(U32.is_lt(6,count),output,(offset + 6 : U32),Bool.pick(F32,biased,(Tile.lane6(accumulator) + bias : F32),Tile.lane6(accumulator)))
87 store_lane(U32.is_lt(7,count),output,(offset + 7 : U32),Bool.pick(F32,biased,(Tile.lane7(accumulator) + bias : F32),Tile.lane7(accumulator)))
88
89def write_bias(output: Array<F32>,offset: U32,accumulator: Tile.Tile8,count: U32,read: Array<F32> & F32) -> Array<F32> & Array<F32>:
90 (bias,value) = read
91 (bias,store_tail(True{},count,output,offset,value,accumulator))
92
93def write_unbiased(bias: Array<F32>,output: Array<F32>,offset: U32,accumulator: Tile.Tile8,count: U32) -> Array<F32> & Array<F32>:
94 (bias,store_tail(False{},count,output,offset,0.0,accumulator))
95
96def write_values(+biased: Bool,bias: Array<F32>,output: Array<F32>,offset: U32,accumulator: Tile.Tile8,count: U32,row: U32) -> Array<F32> & Array<F32>:
97 match biased:
98 case True{}: write_bias(output,offset,accumulator,count,Array.get(F32,bias,row))
99 case False{}: write_unbiased(bias,output,offset,accumulator,count)
100
101def write_row(+biased: Bool,valid: Bool,+row: U32,+position: U32,+columns: U32,accumulator: Tile.Tile8,
102 owners: Array<F32> & Array<F32>) -> Array<F32> & Array<F32>:
103 match valid:
104 case False{}: owners
105 case True{}:
106 (bias,output) = owners
107 write_values(biased,bias,output,(row * columns + position : U32),accumulator,U32.min(8,(columns - position : U32)),row)
108
109def write(+biased: Bool,context: Context,+accumulators: Rows.Rows<Tile.Tile8>,bias: Array<F32>,output: Array<F32>) -> Array<F32> & Array<F32>:
110 Context{+row,+position,+rows,+columns,input_base,weight_base} = context
111 owners = write_row(biased,U32.is_lt(row,rows),row,position,columns,Rows.first(~Tile.Tile8,accumulators),(bias,output))
112 owners = write_row(biased,U32.is_lt((row + 1 : U32),rows),(row + 1 : U32),position,columns,Rows.second(~Tile.Tile8,accumulators),owners)
113 owners = write_row(biased,U32.is_lt((row + 2 : U32),rows),(row + 2 : U32),position,columns,Rows.third(~Tile.Tile8,accumulators),owners)
114 write_row(biased,U32.is_lt((row + 3 : U32),rows),(row + 3 : U32),position,columns,Rows.fourth(~Tile.Tile8,accumulators),owners)
115
116def after_write(input: Array<Tile.Tile8>,weights: Array<Tile.Tile8>,written: Array<F32> & Array<F32>) -> Buffers:
117 (bias,output) = written
118 Buffers{input,weights,bias,output}
119
120# The loop's final step stores its accumulators; bias is added once.
121def finish(~biased: Bool,state: Loop) -> Buffers:
122 Loop{input,weights,bias,output,index,+context,accumulators} = state
123 after_write(input,weights,write(biased,context,accumulators,bias,output))
124
125def reduce(~biased: Bool,upper: Bool,count: Nat,state: Loop) -> Buffers:
126 match upper:
127 case False{}: Iteration.run_then(~Loop,~Buffers,~step(~False{}),~finish(~biased),count,state)
128 case True{}: Iteration.run_then(~Loop,~Buffers,~step(~True{}),~finish(~biased),count,state)
129
130# Every output folds k in order from +0, then adds its bias once.
131def execute_with(~biased: Bool,job: Job,buffers: Buffers,accumulators: Rows.Rows<Tile.Tile8>) -> Buffers:
132 Job{context,upper,count} = job
133 Buffers{input,weights,bias,output} = buffers
134 reduce(~biased,upper,count,Loop{input,weights,bias,output,0,context,accumulators})
135
136def execute(~biased: Bool,job: Job,buffers: Buffers) -> Buffers:
137 execute_with(~biased,job,buffers,Rows.filled(~Tile.Tile8,Tile.filled(0.0)))
138
139# Continue an ordered reduction from four rows already stored in the output.
140def seeded_result(input: Array<Tile.Tile8>,weights: Array<Tile.Tile8>,bias: Array<F32>,+job: Job,
141 a: Tile.Tile8,b: Tile.Tile8,c: Tile.Tile8,read: Array<F32> & Tile.Tile8) -> Buffers:
142 (output,d) = read
143 execute_with(~False{},job,Buffers{input,weights,bias,output},Rows.Rows{a,b,c,d})
144
145def seeded_third(input: Array<Tile.Tile8>,weights: Array<Tile.Tile8>,bias: Array<F32>,+job: Job,+address: U32,+columns: U32,
146 a: Tile.Tile8,b: Tile.Tile8,read: Array<F32> & Tile.Tile8) -> Buffers:
147 (output,c) = read
148 seeded_result(input,weights,bias,job,a,b,c,Lanes.gather(~Unit,~(context => value => value),output,(address + columns : U32),Unit{}))
149
150def seeded_second(input: Array<Tile.Tile8>,weights: Array<Tile.Tile8>,bias: Array<F32>,+job: Job,+address: U32,+columns: U32,
151 a: Tile.Tile8,read: Array<F32> & Tile.Tile8) -> Buffers:
152 (output,b) = read
153 seeded_third(input,weights,bias,job,(address + columns : U32),columns,a,b,Lanes.gather(~Unit,~(context => value => value),output,(address + columns : U32),Unit{}))
154
155def seeded_first(input: Array<Tile.Tile8>,weights: Array<Tile.Tile8>,bias: Array<F32>,+job: Job,+address: U32,+columns: U32,
156 read: Array<F32> & Tile.Tile8) -> Buffers:
157 (output,a) = read
158 seeded_second(input,weights,bias,job,(address + columns : U32),columns,a,Lanes.gather(~Unit,~(context => value => value),output,(address + columns : U32),Unit{}))
159
160def execute_seeded(+job: Job,buffers: Buffers) -> Buffers:
161 Job{Context{+row,+position,rows,+columns,input_base,weight_base},upper,count} = job
162 Buffers{input,weights,bias,output} = buffers
163 +address = (row * columns + position : U32)
164 seeded_first(input,weights,bias,job,address,columns,Lanes.gather(~Unit,~(context => value => value),output,address,Unit{}))
165
166def seeded_job(index: U32,layout: Layout,buffers: Buffers) -> Buffers:
167 execute_seeded(job(layout,index),buffers)
168
169def continue_keeping(+rows: U32,input: Array<Tile.Tile8>,weights: Array<Tile.Tile8>,bias: Array<F32>,output: Array<F32>,+columns: U32,inner: U32) -> Buffers:
170 Traversal.indexed_fold(~Buffers,~Layout,~seeded_job,U32.to_nat(job_count(rows,columns)),0,Layout{rows,columns,inner,0,columns},Buffers{input,weights,bias,output})
171
172def step_job(~biased: Bool,index: U32,layout: Layout,buffers: Buffers) -> Buffers:
173 execute(~biased,job(layout,index),buffers)
174
175def result(buffers: Buffers) -> Array<F32>:
176 Buffers{input,weights,bias,output} = buffers
177 output
178
179# Every job in order. Only the output changes, so the caller gets its other
180# three buffers back.
181def run_keeping(~biased: Bool,+rows: U32,input: Array<Tile.Tile8>,weights: Array<Tile.Tile8>,bias: Array<F32>,output: Array<F32>,+columns: U32,inner: U32,start: U32,output_columns: U32) -> Buffers:
182 Traversal.indexed_fold(~Buffers,~Layout,~step_job(~biased),U32.to_nat(job_count(rows,columns)),0,Layout{rows,columns,inner,start,output_columns},Buffers{input,weights,bias,output})
183
184def run_at(~biased: Bool,+rows: U32,input: Array<Tile.Tile8>,weights: Array<Tile.Tile8>,bias: Array<F32>,output: Array<F32>,+columns: U32,inner: U32,start: U32,output_columns: U32) -> Array<F32>:
185 result(run_keeping(~biased,rows,input,weights,bias,output,columns,inner,start,output_columns))
186
187def run(~biased: Bool,+rows: U32,input: Array<Tile.Tile8>,weights: Array<Tile.Tile8>,bias: Array<F32>,output: Array<F32>,+columns: U32,inner: U32) -> Array<F32>:
188 run_at(~biased,rows,input,weights,bias,output,columns,inner,0,columns)
189
190def kept_input(buffers: Buffers) -> Array<Tile.Tile8>:
191 Buffers{input,weights,bias,output} = buffers
192 input
193
194def kept_weights(buffers: Buffers) -> Array<Tile.Tile8>:
195 Buffers{input,weights,bias,output} = buffers
196 weights
197
198def kept_bias(buffers: Buffers) -> Array<F32>:
199 Buffers{input,weights,bias,output} = buffers
200 bias