kernel_pipeline.bend checks

raw source on the hub · import stelliferous@0.0.2.0/kernel_pipeline.bend as Kernel_pipeline

Each leaf owns its input, packs it, and computes before returning. The caller validates dimensions and chooses depth; the runtime schedules parallel leaves.

16 imports

import Base
import ./kernel_gemm_buffer.bend as GemmBuffer
import ./kernel_tile.bend as Tile
import ./kernel_spatial_packing.bend as Packing
import ./kernel_packing_buffer.bend as PackingBuffer
import ./kernel_gather_buffer.bend as GatherBuffer
import ./storage_buffer.bend as Storage
import ./storage_clone.bend as Clone
import ./tensor.bend as Tensors
import ./traversal_partition.bend as Partition
import ./storage_copy.bend as Copy
import ./kernel_window.bend as KernelWindow
import ./traversal_array.bend as Traversal
import ./math_activation.bend as Activations
import ./proofs/traversal_certificate.bend as TraversalCertificate
import ./tensor_operations.bend as Operations

Types

type Work source · line 20 · raw

Data

Work

@rows:U32 -> @height:U32 -> @width:U32 -> @kernel:U32 -> @stride:U32 -> @padding:U32 -> @output_width:U32 -> @positions:U32 -> @inner:U32 -> Work

type Task source · line 24 · raw

Data

A partition node's work and the activation its leaves apply.

Task

@work:Work -> @activation:0x9be3b13bf249759bc81e1958bcd1a4c0/math_activation.Activation -> Task

type InputWindow source · line 66 · raw

Type

Rows [first, first + rows) of every channel of the image, stored contiguously.

InputWindow

@buffer:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> @first:U32 -> @rows:U32 -> InputWindow

type BlockState source · line 119 · raw

Type

BlockState

@input:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> @tiles:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @weights:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @bias:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> @output:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> BlockState

Definitions

def Outputs source · line 27 · raw

Type

def leaky_relu_value source · line 38 · raw

@slope:F32 -> @value:F32 -> F32

def copy_rows source · line 69 · raw

@remaining:Nat -> @+channel:U32 -> @+count:U32 -> @+source_stride:U32 -> @+source_start:U32 -> @+target_stride:U32 -> @owners:Pair(0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>, 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>) -> Pair(0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>, 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>)

def windows source · line 78 · raw

@+first:U32 -> @+rows:U32 -> @+left_first:U32 -> @+left_rows:U32 -> @owners:Pair(0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>, 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>) -> Pair(InputWindow, InputWindow)

def cloned source · line 82 · raw

@+first:U32 -> @+rows:U32 -> @owners:Pair(0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>, 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>) -> Pair(InputWindow, InputWindow)

def copied source · line 86 · raw

@+channels:U32 -> @+width:U32 -> @+first:U32 -> @+rows:U32 -> @band:0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_window.Band -> @buffer:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> Pair(InputWindow, InputWindow)

def split_when source · line 91 · raw

@copy:Bool -> @+channels:U32 -> @+width:U32 -> @+first:U32 -> @+rows:U32 -> @+band:0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_window.Band -> @buffer:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> Pair(InputWindow, InputWindow)

def split source · line 99 · raw

@work:Work -> @offset:U32 -> @half:U32 -> @window:InputWindow -> Pair(InputWindow, InputWindow)

The right child keeps the parent's rows; the left child copies only its band. Unaligned rows would break whole eight-position blocks, so they are cloned; so is a band that fails its coverage check.

def split_task source · line 108 · raw

@task:Task -> @offset:U32 -> @half:U32 -> @window:InputWindow -> Pair(InputWindow, InputWindow)

def block_tiles source · line 113 · raw

U32

One tile is 32 bytes; 8192 tiles are 256 KB.

def block_positions source · line 116 · raw

@inner:U32 -> U32

def block_kept source · line 122 · raw

@input:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> @kept:0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_gemm_buffer.Kept -> BlockState

def block_packed source · line 126 · raw

@+rows:U32 -> @+positions:U32 -> @inner:U32 -> @+start:U32 -> @+span:U32 -> @weights:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @bias:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> @output:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> @packed:Pair(0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>, 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8>) -> BlockState

def block_step source · line 134 · raw

@work:Work -> @pack_start:U32 -> @output_start:U32 -> @+span:U32 -> @state:BlockState -> BlockState

Packing addresses the input window; output positions are local to the leaf. Each block overwrites tiles from zero and writes directly into the leaf output.

def blocks source · line 140 · raw

@remaining:Nat -> @+work:Work -> @+pack_start:U32 -> @+output_start:U32 -> @+padded:U32 -> @+span:U32 -> @state:BlockState -> BlockState

def gather source · line 172 · raw

@partitions:Outputs -> @output:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> @+positions:U32 -> 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>

def finish_leaf source · line 175 · raw

@identity:Bool -> @input:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> @+rows:U32 -> @offset:U32 -> @span:U32 -> @+positions:U32 -> 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>

def finish source · line 180 · raw

@partitions:Outputs -> @+rows:U32 -> @+positions:U32 -> 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>

def depth_limit source · line 188 · raw

@+padded:U32 -> U32

The deepest split that leaves every leaf at least one eight-position block.

Templates

template activated_value source · line 30 · raw

@-activation:(@_:F32 -> F32) -> @context:Unit -> @value:F32 -> F32

template mapped source · line 33 · raw

@-Context:Data -> @-operation:(@_:Context -> @_:F32 -> F32) -> @+count:U32 -> @+context:Context -> @buffer:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>

template unit_mapped source · line 41 · raw

@-operation:(@_:F32 -> F32) -> @+count:U32 -> @buffer:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>

template unit_lanes source · line 44 · raw

@-operation:(@_:F32 -> F32) -> @+count:U32 -> @buffer:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>

template activated source · line 51 · raw

@-custom:(@_:F32 -> F32) -> @activation:0x9be3b13bf249759bc81e1958bcd1a4c0/math_activation.Activation -> @+count:U32 -> @buffer:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>

One tight pass over each finished leaf while it is still in cache; applying the activation inside the store measured about twice as slow. The built-in is chosen once per leaf, so each pass is specialized; the identity skips it. Sigmoid, SiLU, tanh and GELU pass their values through tiles.

template blocks_finished source · line 147 · raw

@-custom:(@_:F32 -> F32) -> @activation:0x9be3b13bf249759bc81e1958bcd1a4c0/math_activation.Activation -> @count:U32 -> @state:BlockState -> 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>

template blocked source · line 153 · raw

@-custom:(@_:F32 -> F32) -> @+activation:0x9be3b13bf249759bc81e1958bcd1a4c0/math_activation.Activation -> @+work:Work -> @pack_start:U32 -> @input:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> @weights:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @bias:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>

Exactly one reusable tile buffer and one output buffer per leaf. An inner larger than the target still needs one complete eight-position group.

template leaf source · line 161 · raw

@-custom:(@_:F32 -> F32) -> @task:Task -> @offset:U32 -> @span:U32 -> @window:InputWindow -> @weights:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @bias:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> Outputs

template partition source · line 168 · raw

@-custom:(@_:F32 -> F32) -> @depth:Nat -> @task:Task -> @offset:U32 -> @span:U32 -> @input:InputWindow -> @weights:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @bias:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> Outputs

template run_partitioned source · line 191 · raw

@-custom:(@_:F32 -> F32) -> @activation:0x9be3b13bf249759bc81e1958bcd1a4c0/math_activation.Activation -> @depth:U32 -> @work:Work -> @input:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> @weights:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @bias:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>

template run_selected source · line 197 · raw

@-custom:(@_:F32 -> F32) -> @blocked_leaf:Bool -> @activation:0x9be3b13bf249759bc81e1958bcd1a4c0/math_activation.Activation -> @depth:U32 -> @work:Work -> @input:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> @weights:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @bias:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>

template run_packed source · line 202 · raw

@-custom:(@_:F32 -> F32) -> @activation:0x9be3b13bf249759bc81e1958bcd1a4c0/math_activation.Activation -> @+depth:U32 -> @work:Work -> @input:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> @weights:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @bias:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>

template run source · line 205 · raw

@-custom:(@_:F32 -> F32) -> @activation:0x9be3b13bf249759bc81e1958bcd1a4c0/math_activation.Activation -> @depth:U32 -> @+work:Work -> @input:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> @weights:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> @bias:0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>