~/bend-docscommunity

kernel_gemm.bend checks

raw source on the hub · import stelliferous@0.0.2.0/kernel_gemm.bend as Kernel_gemm

5 imports
import Base
import ./kernel_rows.bend as Rows
import ./kernel_tile.bend as Tile
import ./traversal_loop.bend as Iteration
import ./traversal_array.bend as Traversal

Types

type Buffers source · line 7 · raw

Type

type Layout source · line 13 · raw

Data

Row-major output, eight contiguous columns; packed weights are [M/8,K,8]. A job is one half of a weight block, four rows, at eight positions. Only valid rows are stored; the last weight block is zero-padded to eight rows.

type Context source · line 18 · raw

Data

input_base and weight_base address the job's first input group and weight tile; the loop index advances both.

type Job source · line 21 · raw

Data

type Loop source · line 28 · raw

Type

One flat loop state. Every owner, the index and the accumulators are fields of this record, and each helper that passes the accumulators on carries few other values: a nested owner record or a longer helper chain made every k step a real call.

Definitions

def row_units source · line 33 · raw

@rows:U32 -> U32

A job owns four output rows: one half of an eight-row weight block.

def column_groups source · line 36 · raw

@columns:U32 -> U32

def job_count source · line 39 · raw

@rows:U32 -> @columns:U32 -> U32

def job source · line 44 · raw

@layout:Layout -> @index:U32 -> Job

Input-major: one input group stays in cache while every row unit passes over it (the BLIS loop order). Units 2b and 2b+1 share weight block b.

def update source · line 52 · raw

@accumulators:0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_rows.Rows<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @+input:0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8 -> @weights:0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_rows.Rows<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_rows.Rows<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8>

def with_input source · line 55 · raw

@weights:Array<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @bias:Array<F32> -> @output:Array<F32> -> @+index:U32 -> @context:Context -> @accumulators:0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_rows.Rows<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @samples:0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_rows.Rows<F32> -> @read:Pair(Array<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8>, 0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8) -> Loop

def store_lane source · line 73 · raw

@inside:Bool -> @output:Array<F32> -> @offset:U32 -> @value:F32 -> Array<F32>

The last group can contain padding lanes; only real columns are written.

def store_tail source · line 78 · raw

@+count:U32 -> @output:Array<F32> -> @+offset:U32 -> @+bias:F32 -> @+accumulator:0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8 -> Array<F32>

def write_bias source · line 88 · raw

@output:Array<F32> -> @offset:U32 -> @accumulator:0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8 -> @count:U32 -> @read:Pair(Array<F32>, F32) -> Pair(Array<F32>, Array<F32>)

def write_row source · line 92 · raw

@valid:Bool -> @+row:U32 -> @+position:U32 -> @+columns:U32 -> @accumulator:0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8 -> @owners:Pair(Array<F32>, Array<F32>) -> Pair(Array<F32>, Array<F32>)

def write source · line 100 · raw

@context:Context -> @+accumulators:0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_rows.Rows<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @bias:Array<F32> -> @output:Array<F32> -> Pair(Array<F32>, Array<F32>)

def after_write source · line 107 · raw

@input:Array<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @weights:Array<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @written:Pair(Array<F32>, Array<F32>) -> Buffers

def finish source · line 112 · raw

@state:Loop -> Buffers

The loop's final step stores its accumulators; bias is added once.

def reduce source · line 116 · raw

@upper:Bool -> @count:Nat -> @state:Loop -> Buffers

def execute source · line 122 · raw

@job:Job -> @buffers:Buffers -> Buffers

Every output folds k in order from +0, then adds its bias once.

def step_job source · line 127 · raw

@index:U32 -> @layout:Layout -> @buffers:Buffers -> Buffers

def result source · line 130 · raw

@buffers:Buffers -> Array<F32>

def run_keeping source · line 136 · raw

@+rows:U32 -> @input:Array<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @weights:Array<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @bias:Array<F32> -> @output:Array<F32> -> @+columns:U32 -> @inner:U32 -> @start:U32 -> @output_columns:U32 -> Buffers

Every job in order. Only the output changes, so the caller gets its other three buffers back.

def run_at source · line 139 · raw

@+rows:U32 -> @input:Array<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @weights:Array<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @bias:Array<F32> -> @output:Array<F32> -> @+columns:U32 -> @inner:U32 -> @start:U32 -> @output_columns:U32 -> Array<F32>

def run source · line 142 · raw

@+rows:U32 -> @input:Array<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @weights:Array<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @bias:Array<F32> -> @output:Array<F32> -> @+columns:U32 -> @inner:U32 -> Array<F32>

def kept_input source · line 145 · raw

@buffers:Buffers -> Array<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8>

def kept_weights source · line 149 · raw

@buffers:Buffers -> Array<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8>

def kept_bias source · line 153 · raw

@buffers:Buffers -> Array<F32>

Templates

template with_weights source · line 60 · raw

@-upper:Bool -> @input:Array<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @bias:Array<F32> -> @output:Array<F32> -> @+index:U32 -> @context:Context -> @+input_base:U32 -> @accumulators:0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_rows.Rows<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8> -> @read:Pair(Array<0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8>, 0x9be3b13bf249759bc81e1958bcd1a4c0/kernel_tile.Tile8) -> Loop

template step source · line 66 · raw

@-upper:Bool -> @state:Loop -> Loop

One ordered k step: read the weight tile, then the input tile, then MAC.