1 import Base 2 import ./kernel_rows.bend as Rows 3 import ./kernel_tile.bend as Tile 4 import ./traversal_loop.bend as Iteration 5 import ./traversal_array.bend as Traversal 6 7 type Buffers is Type: 8 Buffers{input: Array,weights: Array,bias: Array,output: Array} 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. 13 type 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. 18 type Context is Data: 19 Context{row: U32,position: U32,rows: U32,columns: U32,input_base: U32,weight_base: U32} 20 21 type 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. 28 type Loop is Type: 29 Loop{input: Array,weights: Array,bias: Array,output: Array,index: U32,context: Context, 30 accumulators: Rows.Rows} 31 32 # A job owns four output rows: one half of an eight-row weight block. 33 def row_units(rows: U32) -> U32: 34 ((rows + 3) / 4 : U32) 35 36 def column_groups(columns: U32) -> U32: 37 (columns / 8 : U32) 38 39 def 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. 44 def 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 52 def update(accumulators: Rows.Rows,+input: Tile.Tile8,weights: Rows.Rows) -> Rows.Rows: 53 Rows.map(~Tile.Tile8,~Tile.Tile8,~F32,~Tile.update,accumulators,input,weights) 54 55 def with_input(weights: Array,bias: Array,output: Array,+index: U32,context: Context, 56 accumulators: Rows.Rows,samples: Rows.Rows,read: Array & Tile.Tile8) -> Loop: 57 (input,tile) = read 58 Loop{input,weights,bias,output,(index + 1 : U32),context,update(accumulators,tile,samples)} 59 60 def with_weights(~upper: Bool,input: Array,bias: Array,output: Array,+index: U32,context: Context,+input_base: U32, 61 accumulators: Rows.Rows,read: Array & 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. 66 def 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. 73 def store_lane(inside: Bool,output: Array,offset: U32,value: F32) -> Array: 74 match inside: 75 case True{}: Array.set(F32,output,offset,value) 76 case False{}: output 77 78 def store_tail(+count: U32,output: Array,+offset: U32,+bias: F32,+accumulator: Tile.Tile8) -> Array: 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 88 def write_bias(output: Array,offset: U32,accumulator: Tile.Tile8,count: U32,read: Array & F32) -> Array & Array: 89 (bias,value) = read 90 (bias,store_tail(count,output,offset,value,accumulator)) 91 92 def write_row(valid: Bool,+row: U32,+position: U32,+columns: U32,accumulator: Tile.Tile8, 93 owners: Array & Array) -> Array & Array: 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 100 def write(context: Context,+accumulators: Rows.Rows,bias: Array,output: Array) -> Array & Array: 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 107 def after_write(input: Array,weights: Array,written: Array & Array) -> 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. 112 def 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 116 def 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. 122 def 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 127 def step_job(index: U32,layout: Layout,buffers: Buffers) -> Buffers: 128 execute(job(layout,index),buffers) 129 130 def result(buffers: Buffers) -> Array: 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. 136 def run_keeping(+rows: U32,input: Array,weights: Array,bias: Array,output: Array,+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 139 def run_at(+rows: U32,input: Array,weights: Array,bias: Array,output: Array,+columns: U32,inner: U32,start: U32,output_columns: U32) -> Array: 140 result(run_keeping(rows,input,weights,bias,output,columns,inner,start,output_columns)) 141 142 def run(+rows: U32,input: Array,weights: Array,bias: Array,output: Array,+columns: U32,inner: U32) -> Array: 143 run_at(rows,input,weights,bias,output,columns,inner,0,columns) 144 145 def kept_input(buffers: Buffers) -> Array: 146 Buffers{input,weights,bias,output} = buffers 147 input 148 149 def kept_weights(buffers: Buffers) -> Array: 150 Buffers{input,weights,bias,output} = buffers 151 weights 152 153 def kept_bias(buffers: Buffers) -> Array: 154 Buffers{input,weights,bias,output} = buffers 155 bias