TorchLean API

NN.Tensor.Internal.Elab.Einsum.Tiling.Semantics

Width-polymorphic semantics for einsum output tiles #

An output tile is a finite family of independent contraction accumulators. This module describes that family for an arbitrary number of lanes. Concrete code generators may keep selected widths in separate scalar registers, but their correctness proofs reduce to the definitions and theorems here.

No theorem in this module selects a hardware width or changes scalar evaluation order.

@[inline]
def TorchLean.Tensor.Internal.Elab.Impl.updateTile {α : Type u} [Add α] {lanes : } (totals : Vector α lanes) (values : Fin lanesα) :
Vector α lanes

Add one value to every lane of a contraction tile.

The lane family is a function rather than another vector so generated code can expose each selected lane value without allocating an intermediate container.

Instances For
    theorem TorchLean.Tensor.Internal.Elab.Impl.updateTile_get {α : Type u} [Add α] {lanes : } (totals : Vector α lanes) (values : Fin lanesα) (lane : Fin lanes) :
    (updateTile totals values)[lane] = totals[lane] + values lane

    Projecting a tile update gives the corresponding scalar update.

    theorem TorchLean.Tensor.Internal.Elab.Impl.ofFn_congr {α : Type u} {lanes : } {values reference : Fin lanesα} (h : values = reference) :
    Vector.ofFn values = Vector.ofFn reference

    Equal lane families determine equal function-backed vectors.

    theorem TorchLean.Tensor.Internal.Elab.Impl.coordinateFoldl_add_eq_coordinateSum {α : Type u} [Add α] [OfNat α 0] (shape : Shape) (values : Coord shapeα) (initial : α) :
    coordinateFoldl shape (fun (total : α) (coordinate : Coord shape) => total + values coordinate) initial = Semantics.coordinateSum shape values initial

    The additive coordinate fold is the executable presentation of Semantics.coordinateSum.

    theorem TorchLean.Tensor.Internal.Elab.Impl.coordinateFoldl_vector {α : Type u} [Add α] [OfNat α 0] {lanes : } (shape : Shape) (update : Vector α lanesCoord shapeVector α lanes) (value : Fin lanesCoord shapeα) (initial : Vector α lanes) (hUpdate : ∀ (totals : Vector α lanes) (coordinate : Coord shape) (selectedLane : Fin lanes), (update totals coordinate)[selectedLane] = totals[selectedLane] + value selectedLane coordinate) :
    coordinateFoldl shape update initial = Vector.ofFn fun (lane : Fin lanes) => Semantics.coordinateSum shape (value lane) initial[lane]

    A vector-valued coordinate fold is the vector of its scalar lane sums when every update is pointwise addition.

    theorem TorchLean.Tensor.Internal.Elab.Impl.coordinateFoldl_updateTile {α : Type u} [Add α] [OfNat α 0] {lanes : } (shape : Shape) (values : Fin lanesCoord shapeα) (initial : α) :
    coordinateFoldl shape (fun (totals : Vector α lanes) (coordinate : Coord shape) => updateTile totals fun (lane : Fin lanes) => values lane coordinate) (Vector.replicate lanes initial) = Vector.ofFn fun (lane : Fin lanes) => Semantics.coordinateSum shape (values lane) initial

    An arbitrary-width contraction tile equals the vector of its scalar contractions. Every lane therefore retains the scalar fold's original coordinate order.

    theorem TorchLean.Tensor.Internal.Elab.Impl.coordinateFoldl_updateTile_of_eq {α : Type u} [Add α] [OfNat α 0] {lanes : } (shape : Shape) (update : Vector α lanesCoord shapeVector α lanes) (values : Fin lanesCoord shapeα) (initial : α) (hUpdate : ∀ (totals : Vector α lanes) (coordinate : Coord shape), update totals coordinate = updateTile totals fun (lane : Fin lanes) => values lane coordinate) :
    coordinateFoldl shape update (Vector.replicate lanes initial) = Vector.ofFn fun (lane : Fin lanes) => Semantics.coordinateSum shape (values lane) initial

    Any concrete tile update that agrees pointwise with updateTile has the same arbitrary-width contraction semantics.

    Scalar-register kernels use this theorem as their only semantic boundary. Their implementation-specific proof need only identify one update step.