kernel_strided_store.bend
stelliferous@0.0.3.0 · 30 lines · 0 definitions
1# Four rows of eight accumulator lanes stored at strided addresses. Each value
2# is the fold of its products from +0, then +0.0.
3import Base
4import stelliferous@0.0.3.0/kernel_tile.bend as Tile
5import stelliferous@0.0.3.0/kernel_rows.bend as Rows
6import stelliferous@0.0.3.0/kernel_gemm.bend as Gemm
7
8# Row r, lane i goes to target + r * plane + i * step, for r < rows and i < count.
9type Store is Data:
10 Store{target: U32,step: U32,plane: U32,rows: U32,count: U32}
11
12def row(valid: Bool,output: Array<F32>,+target: U32,+step: U32,+count: U32,+accumulator: Tile.Tile8) -> Array<F32>:
13 match valid:
14 case False{}: output
15 case True{}:
16 output = Gemm.store_lane(U32.is_lt(0,count),output,target,(Tile.lane0(accumulator) + 0.0 : F32))
17 output = Gemm.store_lane(U32.is_lt(1,count),output,(target + step : U32),(Tile.lane1(accumulator) + 0.0 : F32))
18 output = Gemm.store_lane(U32.is_lt(2,count),output,(target + 2 * step : U32),(Tile.lane2(accumulator) + 0.0 : F32))
19 output = Gemm.store_lane(U32.is_lt(3,count),output,(target + 3 * step : U32),(Tile.lane3(accumulator) + 0.0 : F32))
20 output = Gemm.store_lane(U32.is_lt(4,count),output,(target + 4 * step : U32),(Tile.lane4(accumulator) + 0.0 : F32))
21 output = Gemm.store_lane(U32.is_lt(5,count),output,(target + 5 * step : U32),(Tile.lane5(accumulator) + 0.0 : F32))
22 output = Gemm.store_lane(U32.is_lt(6,count),output,(target + 6 * step : U32),(Tile.lane6(accumulator) + 0.0 : F32))
23 Gemm.store_lane(U32.is_lt(7,count),output,(target + 7 * step : U32),(Tile.lane7(accumulator) + 0.0 : F32))
24
25def block(+store: Store,+accumulators: Rows.Rows<Tile.Tile8>,output: Array<F32>) -> Array<F32>:
26 Store{+target,+step,+plane,+rows,+count} = store
27 output = row(U32.is_lt(0,rows),output,target,step,count,Rows.first(~Tile.Tile8,accumulators))
28 output = row(U32.is_lt(1,rows),output,(target + plane : U32),step,count,Rows.second(~Tile.Tile8,accumulators))
29 output = row(U32.is_lt(2,rows),output,(target + 2 * plane : U32),step,count,Rows.third(~Tile.Tile8,accumulators))
30 row(U32.is_lt(3,rows),output,(target + 3 * plane : U32),step,count,Rows.fourth(~Tile.Tile8,accumulators))