Four-register output tiling

This module implements one concrete scalar-register lowering selected by static cost analysis. Its update step refines the arbitrary-width semantics from Tiling.Semantics; the width is an implementation choice rather than a semantic restriction.

Every lane observes contracted coordinates in the same row-major order as the scalar lowering, so tiling requires no reassociation, commutativity, or distributivity argument.

def TorchLean.Tensor.Internal.Elab.Impl.selectTile4 {α : Type u} (value0 value1 value2 value3 : α) : Fin 4 → α

Select one of four values by a statically bounded lane index.

theorem TorchLean.Tensor.Internal.Elab.Impl.coordinateSum_selectTile4 {α : Type u} [Add α] [OfNat α 0] (shape : Shape) (value0 value1 value2 value3 : Coord shape → α) (initial : α) : (fun (lane : Fin 4) => Semantics.coordinateSum shape (fun (coordinate : Coord shape) => selectTile4 (value0 coordinate) (value1 coordinate) (value2 coordinate) (value3 coordinate) lane) initial) = selectTile4 (Semantics.coordinateSum shape value0 initial) (Semantics.coordinateSum shape value1 initial) (Semantics.coordinateSum shape value2 initial) (Semantics.coordinateSum shape value3 initial)

Selecting a concrete four-lane contraction after summing each lane agrees with summing the selected lane. This representation lemma preserves the coordinate order of Semantics.coordinateSum.

@[inline]

def TorchLean.Tensor.Internal.Elab.Impl.updateTile4 {α : Type u} [Add α] (totals : Vector α 4) (value0 value1 value2 value3 : α) : Vector α 4

Update the proof-facing state of a four-lane tile.

For concrete contractions, compileCoordinateFold recognizes this exact update and first emits nativeFinSum4, whose four totals are separate loop arguments. The surrounding output pass fuses an immediately consumed result into nativeFinSum4Push.

@[inline]

def TorchLean.Tensor.Internal.Elab.Impl.pushTile4 {α : Type u} [storage : Storage α] (output : Storage.Buffer α) (values : Vector α 4) : Storage.Buffer α

Append one completed four-lane tile to a row-major output buffer.

theorem TorchLean.Tensor.Internal.Elab.Impl.ofFn_selectTile4_congr {α : Type u} (value0 value1 value2 value3 reference0 reference1 reference2 reference3 : α) (h0 : value0 = reference0) (h1 : value1 = reference1) (h2 : value2 = reference2) (h3 : value3 = reference3) : Vector.ofFn (selectTile4 value0 value1 value2 value3) = Vector.ofFn (selectTile4 reference0 reference1 reference2 reference3)

Pointwise equal lane values determine equal four-lane vectors.

theorem TorchLean.Tensor.Internal.Elab.Impl.pushTile4_ofFn {α : Type u} [storage : Storage α] (output : Storage.Buffer α) (value0 value1 value2 value3 : α) : pushTile4 output (Vector.ofFn (selectTile4 value0 value1 value2 value3)) = Storage.push (Storage.push (Storage.push (Storage.push output value0) value1) value2) value3

Appending a four-lane function vector emits its values in lane order.

theorem TorchLean.Tensor.Internal.Elab.Impl.updateTile4_get {α : Type u} [Add α] (totals : Vector α 4) (value0 value1 value2 value3 : α) (lane : Fin 4) : (updateTile4 totals value0 value1 value2 value3)[lane] = totals[lane] + selectTile4 value0 value1 value2 value3 lane

Projecting a four-lane update gives the corresponding scalar update.

theorem TorchLean.Tensor.Internal.Elab.Impl.updateTile4_eq_updateTile {α : Type u} [Add α] (totals : Vector α 4) (value0 value1 value2 value3 : α) : updateTile4 totals value0 value1 value2 value3 = updateTile totals (selectTile4 value0 value1 value2 value3)

The four-register update implements the width-polymorphic tile update.

@[irreducible, specialize #[]]

def TorchLean.Tensor.Internal.Elab.Impl.nativeFinSum4Loop {α : Type u} [Add α] (length : ℕ) (value0 value1 value2 value3 : (index : USize) → index.toNat < length → α) (bound : USize) (hBound : bound.toNat = length) (index : USize) (total0 total1 total2 total3 : α) : Vector α 4

Run a native finite sum while carrying four lane totals as separate scalar arguments.

Each recursive step advances one contraction coordinate. This keeps the generated callback small enough for Lean's native compiler to inline it into the loop, matching the eight-lane implementation.

@[inline]

def TorchLean.Tensor.Internal.Elab.Impl.nativeFinSum4 {α : Type u} [Add α] (length : ℕ) (bound : USize) (hBound : bound.toNat = length) (value0 value1 value2 value3 : (index : USize) → index.toNat < length → α) (initial : Vector α 4) : Vector α 4

Sum four lanes over a native finite interval without vector operations inside the loop.

@[irreducible, specialize #[]]

def TorchLean.Tensor.Internal.Elab.Impl.nativeFinSum4PushLoop {α : Type u} [storage : Storage α] [Add α] (length : ℕ) (value0 value1 value2 value3 : (index : USize) → index.toNat < length → α) (bound : USize) (hBound : bound.toNat = length) (output : Storage.Buffer α) (index : USize) (total0 total1 total2 total3 : α) : Storage.Buffer α

Run four scalar accumulators and append their final values directly to an existing output buffer.

@[inline]

def TorchLean.Tensor.Internal.Elab.Impl.nativeFinSum4Push {α : Type u} [storage : Storage α] [Add α] (length : ℕ) (bound : USize) (hBound : bound.toNat = length) (value0 value1 value2 value3 : (index : USize) → index.toNat < length → α) (output : Storage.Buffer α) (initial : α) : Storage.Buffer α

Sum four lanes and append them without materializing the result vector.

theorem TorchLean.Tensor.Internal.Elab.Impl.nativeFinSum4_eq_nativeFinFoldl {α : Type u} [Add α] (length : ℕ) (bound : USize) (hBound : bound.toNat = length) (value0 value1 value2 value3 : (index : USize) → index.toNat < length → α) (initial : Vector α 4) : nativeFinSum4 length bound hBound value0 value1 value2 value3 initial = nativeFinFoldl length bound hBound (fun (totals : Vector α 4) (index : USize) (hIndex : index.toNat < length) => updateTile4 totals (value0 index hIndex) (value1 index hIndex) (value2 index hIndex) (value3 index hIndex)) initial

The four-scalar native loop is exactly the ordinary native fold using updateTile4. Every lane therefore retains its original reduction order.

theorem TorchLean.Tensor.Internal.Elab.Impl.nativeFinSum4Push_eq_pushTile4 {α : Type u} [storage : Storage α] [Add α] (length : ℕ) (bound : USize) (hBound : bound.toNat = length) (value0 value1 value2 value3 : (index : USize) → index.toNat < length → α) (output : Storage.Buffer α) (initial : α) : nativeFinSum4Push length bound hBound value0 value1 value2 value3 output initial = pushTile4 output (nativeFinSum4 length bound hBound value0 value1 value2 value3 (Vector.replicate 4 initial))

Directly appending the scalar loop totals is exactly pushTile4 applied to the vector-valued loop result.

@[irreducible, specialize #[]]

def TorchLean.Tensor.Internal.Elab.Impl.nativeFinSum4FinalizePushLoop {α : Type u} [storage : Storage α] [Add α] (length : ℕ) (value0 value1 value2 value3 : (index : USize) → index.toNat < length → α) (finalize0 finalize1 finalize2 finalize3 : α → α) (bound : USize) (hBound : bound.toNat = length) (output : Storage.Buffer α) (index : USize) (total0 total1 total2 total3 : α) : Storage.Buffer α

Run four scalar accumulators, apply one terminal function to each total, and append the results directly to an existing output buffer.

The terminal functions are never called inside the contraction loop. This is the execution path used when semiring laws move contraction-invariant factors outside a shared tiled sum.

@[inline]

def TorchLean.Tensor.Internal.Elab.Impl.nativeFinSum4FinalizePush {α : Type u} [storage : Storage α] [Add α] (length : ℕ) (bound : USize) (hBound : bound.toNat = length) (value0 value1 value2 value3 : (index : USize) → index.toNat < length → α) (finalize0 finalize1 finalize2 finalize3 : α → α) (output : Storage.Buffer α) (initial : α) : Storage.Buffer α

Sum four lanes and apply their terminal functions directly at the output buffer boundary.

theorem TorchLean.Tensor.Internal.Elab.Impl.nativeFinSum4FinalizePush_eq_pushTile4 {α : Type u} [storage : Storage α] [Add α] (length : ℕ) (bound : USize) (hBound : bound.toNat = length) (value0 value1 value2 value3 : (index : USize) → index.toNat < length → α) (finalize0 finalize1 finalize2 finalize3 : α → α) (output : Storage.Buffer α) (initial : α) : nativeFinSum4FinalizePush length bound hBound value0 value1 value2 value3 finalize0 finalize1 finalize2 finalize3 output initial = have totals := nativeFinSum4 length bound hBound value0 value1 value2 value3 (Vector.replicate 4 initial); pushTile4 output (Vector.ofFn (selectTile4 (finalize0 totals[0]) (finalize1 totals[1]) (finalize2 totals[2]) (finalize3 totals[3])))

The finalized four-lane append is exactly the vector-valued contraction followed by lane-wise finalization and pushTile4.