kernel_gemm.bend

stelliferous@0.0.2.0 · 155 lines · 30 definitions

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