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.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 α lanes → Coord shape → Vector α lanes) (value : Fin lanes → Coord 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 lanes → Coord 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 α lanes → Coord shape → Vector α lanes) (values : Fin lanes → Coord 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.