def blocked

stelliferous@0.0.2.0 · kernel_pipeline.bend:151

import Base
import stelliferous@0.0.2.0/kernel_tile.bend as Tile
import stelliferous@0.0.2.0/storage_buffer.bend as Storage
import stelliferous@0.0.2.0/tensor.bend as Tensors
import stelliferous@0.0.2.0/kernel_window.bend as KernelWindow
import stelliferous@0.0.2.0/math_activation.bend as Activations

151# Exactly one reusable tile buffer and one output buffer per leaf.
152# An inner larger than the target still needs one complete eight-position group.
153def blocked(~custom: F32 -> F32,+activation: Activations.Activation,+work: Work,pack_start: U32,input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
154  Work{+rows,h,w,k,s,p,ow,+positions,+inner} = work
155  +padded = KernelWindow.padded(positions)
156  +span = block_positions(inner)
157  blocks_finished(~custom,activation,(rows * positions : U32),
158    blocks(U32.to_nat(((padded + span - 1) / span : U32)),work,pack_start,0,padded,span,
159      BlockState{input,Tensors.buffer(~Tile.Tile8,Tile.filled(0.0),U32.max(block_tiles(),inner)),weights,bias,Tensors.buffer(~F32,0.0,(rows * positions : U32))}))