TorchLean API

NN.Tensor.Internal.Elab.Einsum.Tiling.Width8

Eight-register output tiling #

This module implements the wider concrete scalar-register lowering selected for large contractions. 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.selectTile8 {α : Type u} (value0 value1 value2 value3 value4 value5 value6 value7 : α) :
Fin 8α

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

Instances For
    theorem TorchLean.Tensor.Internal.Elab.Impl.coordinateSum_selectTile8 {α : Type u} [Add α] [OfNat α 0] (shape : Shape) (value0 value1 value2 value3 value4 value5 value6 value7 : Coord shapeα) (initial : α) :
    (fun (lane : Fin 8) => Semantics.coordinateSum shape (fun (coordinate : Coord shape) => selectTile8 (value0 coordinate) (value1 coordinate) (value2 coordinate) (value3 coordinate) (value4 coordinate) (value5 coordinate) (value6 coordinate) (value7 coordinate) lane) initial) = selectTile8 (Semantics.coordinateSum shape value0 initial) (Semantics.coordinateSum shape value1 initial) (Semantics.coordinateSum shape value2 initial) (Semantics.coordinateSum shape value3 initial) (Semantics.coordinateSum shape value4 initial) (Semantics.coordinateSum shape value5 initial) (Semantics.coordinateSum shape value6 initial) (Semantics.coordinateSum shape value7 initial)

    Selecting a concrete eight-lane contraction after summing each lane agrees with summing the selected lane.

    @[inline]
    def TorchLean.Tensor.Internal.Elab.Impl.updateTile8 {α : Type u} [Add α] (totals : Vector α 8) (value0 value1 value2 value3 value4 value5 value6 value7 : α) :
    Vector α 8

    Update the proof-facing state of an eight-lane tile.

    For large concrete contractions, compileCoordinateFold recognizes this exact update and first emits nativeFinSum8, which carries the eight totals as separate loop arguments. The surrounding output pass fuses an immediately consumed result into nativeFinSum8Push. The scalar-state theorems prove both forms compute this vector update exactly.

    Instances For
      @[inline]
      def TorchLean.Tensor.Internal.Elab.Impl.pushTile8 {α : Type u} [storage : Storage α] (output : Storage.Buffer α) (values : Vector α 8) :

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

      Instances For
        theorem TorchLean.Tensor.Internal.Elab.Impl.ofFn_selectTile8_congr {α : Type u} (value0 value1 value2 value3 value4 value5 value6 value7 reference0 reference1 reference2 reference3 reference4 reference5 reference6 reference7 : α) (h0 : value0 = reference0) (h1 : value1 = reference1) (h2 : value2 = reference2) (h3 : value3 = reference3) (h4 : value4 = reference4) (h5 : value5 = reference5) (h6 : value6 = reference6) (h7 : value7 = reference7) :
        Vector.ofFn (selectTile8 value0 value1 value2 value3 value4 value5 value6 value7) = Vector.ofFn (selectTile8 reference0 reference1 reference2 reference3 reference4 reference5 reference6 reference7)

        Pointwise equal lane values determine equal eight-lane vectors.

        theorem TorchLean.Tensor.Internal.Elab.Impl.pushTile8_ofFn {α : Type u} [storage : Storage α] (output : Storage.Buffer α) (value0 value1 value2 value3 value4 value5 value6 value7 : α) :
        pushTile8 output (Vector.ofFn (selectTile8 value0 value1 value2 value3 value4 value5 value6 value7)) = have output := Storage.push output value0; have output := Storage.push output value1; have output := Storage.push output value2; have output := Storage.push output value3; have output := Storage.push output value4; have output := Storage.push output value5; have output := Storage.push output value6; Storage.push output value7

        Appending an eight-lane function vector emits its values in lane order.

        theorem TorchLean.Tensor.Internal.Elab.Impl.updateTile8_get {α : Type u} [Add α] (totals : Vector α 8) (value0 value1 value2 value3 value4 value5 value6 value7 : α) (lane : Fin 8) :
        (updateTile8 totals value0 value1 value2 value3 value4 value5 value6 value7)[lane] = totals[lane] + selectTile8 value0 value1 value2 value3 value4 value5 value6 value7 lane

        Projecting an eight-lane update gives the corresponding scalar update.

        theorem TorchLean.Tensor.Internal.Elab.Impl.updateTile8_eq_updateTile {α : Type u} [Add α] (totals : Vector α 8) (value0 value1 value2 value3 value4 value5 value6 value7 : α) :
        updateTile8 totals value0 value1 value2 value3 value4 value5 value6 value7 = updateTile totals (selectTile8 value0 value1 value2 value3 value4 value5 value6 value7)

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

        @[irreducible, specialize #[]]
        def TorchLean.Tensor.Internal.Elab.Impl.nativeFinSum8Loop {α : Type u} [Add α] (length : ) (value0 value1 value2 value3 value4 value5 value6 value7 : (index : USize) → index.toNat < lengthα) (bound : USize) (hBound : bound.toNat = length) (index : USize) (total0 total1 total2 total3 total4 total5 total6 total7 : α) :
        Vector α 8

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

        The compiler selects this loop when eight neighboring outputs are available and shared coordinate work amortizes the additional live state. Both concrete tile widths advance one contraction coordinate per recursive step; the eight-lane loop keeps twice as many independent accumulators live.

        Instances For
          @[inline]
          def TorchLean.Tensor.Internal.Elab.Impl.nativeFinSum8 {α : Type u} [Add α] (length : ) (bound : USize) (hBound : bound.toNat = length) (value0 value1 value2 value3 value4 value5 value6 value7 : (index : USize) → index.toNat < lengthα) (initial : Vector α 8) :
          Vector α 8

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

          Instances For
            @[irreducible, specialize #[]]
            def TorchLean.Tensor.Internal.Elab.Impl.nativeFinSum8PushLoop {α : Type u} [storage : Storage α] [Add α] (length : ) (value0 value1 value2 value3 value4 value5 value6 value7 : (index : USize) → index.toNat < lengthα) (bound : USize) (hBound : bound.toNat = length) (output : Storage.Buffer α) (index : USize) (total0 total1 total2 total3 total4 total5 total6 total7 : α) :

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

            The terminal pushes avoid constructing a temporary result vector between the contraction loop and the surrounding output loop.

            Instances For
              @[inline]
              def TorchLean.Tensor.Internal.Elab.Impl.nativeFinSum8Push {α : Type u} [storage : Storage α] [Add α] (length : ) (bound : USize) (hBound : bound.toNat = length) (value0 value1 value2 value3 value4 value5 value6 value7 : (index : USize) → index.toNat < lengthα) (output : Storage.Buffer α) (initial : α) :

              Sum eight lanes and append them without materializing the intermediate vector.

              Instances For
                theorem TorchLean.Tensor.Internal.Elab.Impl.nativeFinSum8_eq_nativeFinFoldl {α : Type u} [Add α] (length : ) (bound : USize) (hBound : bound.toNat = length) (value0 value1 value2 value3 value4 value5 value6 value7 : (index : USize) → index.toNat < lengthα) (initial : Vector α 8) :
                nativeFinSum8 length bound hBound value0 value1 value2 value3 value4 value5 value6 value7 initial = nativeFinFoldl length bound hBound (fun (totals : Vector α 8) (index : USize) (hIndex : index.toNat < length) => updateTile8 totals (value0 index hIndex) (value1 index hIndex) (value2 index hIndex) (value3 index hIndex) (value4 index hIndex) (value5 index hIndex) (value6 index hIndex) (value7 index hIndex)) initial

                The eight-scalar native loop is exactly the ordinary native fold using updateTile8. In particular, every lane retains its original reduction order.

                theorem TorchLean.Tensor.Internal.Elab.Impl.nativeFinSum8Push_eq_pushTile8 {α : Type u} [storage : Storage α] [Add α] (length : ) (bound : USize) (hBound : bound.toNat = length) (value0 value1 value2 value3 value4 value5 value6 value7 : (index : USize) → index.toNat < lengthα) (output : Storage.Buffer α) (initial : α) :
                nativeFinSum8Push length bound hBound value0 value1 value2 value3 value4 value5 value6 value7 output initial = pushTile8 output (nativeFinSum8 length bound hBound value0 value1 value2 value3 value4 value5 value6 value7 (Vector.replicate 8 initial))

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

                @[irreducible, specialize #[]]
                def TorchLean.Tensor.Internal.Elab.Impl.nativeFinSum8FinalizePushLoop {α : Type u} [storage : Storage α] [Add α] (length : ) (value0 value1 value2 value3 value4 value5 value6 value7 : (index : USize) → index.toNat < lengthα) (finalize0 finalize1 finalize2 finalize3 finalize4 finalize5 finalize6 finalize7 : αα) (bound : USize) (hBound : bound.toNat = length) (output : Storage.Buffer α) (index : USize) (total0 total1 total2 total3 total4 total5 total6 total7 : α) :

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

                The terminal functions run only after the contraction has completed. This keeps factored semiring operands out of the hot loop without introducing an intermediate result vector.

                Instances For
                  @[inline]
                  def TorchLean.Tensor.Internal.Elab.Impl.nativeFinSum8FinalizePush {α : Type u} [storage : Storage α] [Add α] (length : ) (bound : USize) (hBound : bound.toNat = length) (value0 value1 value2 value3 value4 value5 value6 value7 : (index : USize) → index.toNat < lengthα) (finalize0 finalize1 finalize2 finalize3 finalize4 finalize5 finalize6 finalize7 : αα) (output : Storage.Buffer α) (initial : α) :

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

                  Instances For
                    theorem TorchLean.Tensor.Internal.Elab.Impl.nativeFinSum8FinalizePush_eq_pushTile8 {α : Type u} [storage : Storage α] [Add α] (length : ) (bound : USize) (hBound : bound.toNat = length) (value0 value1 value2 value3 value4 value5 value6 value7 : (index : USize) → index.toNat < lengthα) (finalize0 finalize1 finalize2 finalize3 finalize4 finalize5 finalize6 finalize7 : αα) (output : Storage.Buffer α) (initial : α) :
                    nativeFinSum8FinalizePush length bound hBound value0 value1 value2 value3 value4 value5 value6 value7 finalize0 finalize1 finalize2 finalize3 finalize4 finalize5 finalize6 finalize7 output initial = have totals := nativeFinSum8 length bound hBound value0 value1 value2 value3 value4 value5 value6 value7 (Vector.replicate 8 initial); pushTile8 output (Vector.ofFn (selectTile8 (finalize0 totals[0]) (finalize1 totals[1]) (finalize2 totals[2]) (finalize3 totals[3]) (finalize4 totals[4]) (finalize5 totals[5]) (finalize6 totals[6]) (finalize7 totals[7])))

                    Applying terminal functions at the end of the eight-scalar loop is equivalent to mapping those functions over the vector-valued loop result before append.