TorchLean API

NN.Tensor.Internal.Elab.Native.Reduce

Certified native reduction loops #

Concrete reductions use one native output loop and one native fiber loop. The pointwise reader certificate connects direct source-buffer reads to the existing row-major reduction semantics, so scalar operation order is unchanged.

@[inline]
def TorchLean.Tensor.Internal.Elab.Impl.nativeReduceFoldTensorFromFlat {α : Type u} {β : Type v} {γ : Type w} [Storage α] [Storage γ] (step : β → α → β) (initial : β) (finish : β → ℕ → γ) (checked : Check.CheckedTransform) (_hKind : checked.value.normalized.kind = Check.TransformKind.reduce) (inputTensor : checked.InputTensor α) (read : Fin checked.value.normalized.input.size → α) (_hRead : ∀ (inputIndex : Fin checked.value.normalized.input.size), read inputIndex = Rep.getFlat inputTensor inputIndex) (outputBound : USize) (hOutputBound : outputBound.toNat = checked.value.output.size) (fiberBound : USize) (hFiberBound : fiberBound.toNat = checked.reductionShape.size) (nativeRead : (outputIndex : USize) → outputIndex.toNat < checked.value.output.size → (fiberIndex : USize) → fiberIndex.toNat < checked.reductionShape.size → α) (hNativeRead : ∀ (outputIndex : USize) (hOutputIndex : outputIndex.toNat < checked.value.output.size) (fiberIndex : USize) (hFiberIndex : fiberIndex.toNat < checked.reductionShape.size), nativeRead outputIndex hOutputIndex fiberIndex hFiberIndex = read (Lowering.Reduce.Impl.reductionInputFlatIndex checked (Coord.unlinearize ⟨outputIndex.toNat, hOutputIndex⟩) ⟨fiberIndex.toNat, hFiberIndex⟩)) :
checked.OutputTensor γ

Reduce every fiber with native output and fiber counters.

The executable callback receives both counters as USize. Its pointwise certificate is erased and proves agreement with the existing flat-reader reduction at the corresponding bounded indices.

Instances For
    theorem TorchLean.Tensor.Internal.Elab.Impl.nativeReduceFoldTensorFromFlat_correct {α : Type u} {β : Type v} {γ : Type w} [Storage α] [Storage γ] (step : β → α → β) (initial : β) (finish : β → ℕ → γ) (checked : Check.CheckedTransform) (hKind : checked.value.normalized.kind = Check.TransformKind.reduce) (inputTensor : checked.InputTensor α) (read : Fin checked.value.normalized.input.size → α) (hRead : ∀ (inputIndex : Fin checked.value.normalized.input.size), read inputIndex = Rep.getFlat inputTensor inputIndex) (outputBound : USize) (hOutputBound : outputBound.toNat = checked.value.output.size) (fiberBound : USize) (hFiberBound : fiberBound.toNat = checked.reductionShape.size) (nativeRead : (outputIndex : USize) → outputIndex.toNat < checked.value.output.size → (fiberIndex : USize) → fiberIndex.toNat < checked.reductionShape.size → α) (hNativeRead : ∀ (outputIndex : USize) (hOutputIndex : outputIndex.toNat < checked.value.output.size) (fiberIndex : USize) (hFiberIndex : fiberIndex.toNat < checked.reductionShape.size), nativeRead outputIndex hOutputIndex fiberIndex hFiberIndex = read (Lowering.Reduce.Impl.reductionInputFlatIndex checked (Coord.unlinearize ⟨outputIndex.toNat, hOutputIndex⟩) ⟨fiberIndex.toNat, hFiberIndex⟩)) :
    nativeReduceFoldTensorFromFlat step initial finish checked hKind inputTensor read hRead outputBound hOutputBound fiberBound hFiberBound nativeRead hNativeRead = Lowering.reduceFoldTensorFromFlat step initial finish checked hKind inputTensor read hRead

    The native ordered reduction is exactly the existing flat-reader reduction.

    def TorchLean.Tensor.Internal.Elab.Impl.compileNativeReduceFold? (step initial finish checked hKind inputTensor semanticInputTensor hInputTensor read hRead sourceTensor inputFlatMap : Lean.Expr) :

    Compile an ordered flat-reader reduction to nested native loops.

    The source map may be the identity map of an ordinary tensor or a certified map recovered from a preceding shape-only transform.

    Nonconcrete or nonportable output, fiber, and logical-input sizes return none. Once compilation begins, certificate errors propagate to the caller; this function does not transactionally catch failed optimizations.

    Instances For