From 412ec8b7fe6e3bb062ea8526ad0b6cbee023bde8 Mon Sep 17 00:00:00 2001 From: Nicolas Rouquette Date: Sat, 3 Oct 2026 14:06:57 -0700 Subject: [PATCH 1/2] feat(libtorch): piecewise-linear table lookup composed from float32 primitives Add `floor` to the shared operation list, a gather at device-resident positions with its scatter-add VJP, and `Tape.tableLookup` / `Tape.layeredTableLookup`: clamp-addressed linear interpolation into a tabulated function, composed in Lean from `clamp`, `floor`, `sub`, `mul`, `add`, and the new gather so that every operation rounds once and the result matches `KernelSpec.tableLookupSpec` bit for bit. The backward gives the coordinate its segment's slope, zero where the clamp is active, and the table the two interpolation weights. Tests compare `floor` with `Float32.floor`, `gatherAt` with the CPU tape's `indexSelect`, the lookups with a `Float32` reference at boundary, out-of-range, integral, and NaN coordinates, and the forward and both gradients with the same function recorded on the CPU tape. A repeated `scatterAddAt` must return the same bits under the deterministic control. --- .../Autograd/Engine/LibTorch/Buffer.lean | 4 + .../Autograd/Engine/LibTorch/KernelSpec.lean | 17 + .../Autograd/Engine/LibTorch/Kernels.lean | 20 + NN/Runtime/Autograd/Engine/LibTorch/Ops.lean | 1 + .../Engine/LibTorch/Ops/Elementwise.lean | 6 + .../Engine/LibTorch/Ops/Indexing.lean | 32 ++ .../Engine/LibTorch/Ops/TableLookup.lean | 178 +++++++++ NN/Tests/Runtime/Cuda/Suite.lean | 2 + NN/Tests/Runtime/Cuda/TableLookup.lean | 352 ++++++++++++++++++ csrc/libtorch/README.md | 1 + csrc/libtorch/operations.h | 1 + csrc/libtorch/torchlean.cpp | 34 ++ csrc/libtorch/unavailable.c | 4 + 13 files changed, 652 insertions(+) create mode 100644 NN/Runtime/Autograd/Engine/LibTorch/Ops/TableLookup.lean create mode 100644 NN/Tests/Runtime/Cuda/TableLookup.lean diff --git a/NN/Runtime/Autograd/Engine/LibTorch/Buffer.lean b/NN/Runtime/Autograd/Engine/LibTorch/Buffer.lean index dda85b72..71893d58 100644 --- a/NN/Runtime/Autograd/Engine/LibTorch/Buffer.lean +++ b/NN/Runtime/Autograd/Engine/LibTorch/Buffer.lean @@ -477,6 +477,10 @@ opaque abs (b : @& Buffer) : Buffer @[never_extract, extern "torchlean_cuda_buffer_abs_bwd"] opaque absBwd (x dLdy : @& Buffer) : Buffer +/-- Elementwise floor: the largest integral value not above each element, as `Float32.floor`. -/ +@[never_extract, extern "torchlean_cuda_buffer_floor"] +opaque floor (b : @& Buffer) : Buffer + /-- Elementwise `sqrt (max x 0)`, matching `Tensor.sqrtSpec`. Negative inputs and either signed zero return positive zero. NaN inputs remain NaN; the clamp diff --git a/NN/Runtime/Autograd/Engine/LibTorch/KernelSpec.lean b/NN/Runtime/Autograd/Engine/LibTorch/KernelSpec.lean index 65c2acc0..bdfefc55 100644 --- a/NN/Runtime/Autograd/Engine/LibTorch/KernelSpec.lean +++ b/NN/Runtime/Autograd/Engine/LibTorch/KernelSpec.lean @@ -298,6 +298,23 @@ def scatterAddSpec {n k : Nat} (x : FlatBuffer n) (values : FlatBuffer k) (fun acc j => if idx j = i then ExecFloat.add acc (values j) else acc) (x i) +/-! ## Table lookup -/ + +/-- +Piecewise-linear table lookup reference: `v₀ + (v₁ - v₀) · w` with one rounding per operation. + +`lower` and `upper` are the sample positions each element interpolates between and `weight` is +its interpolation weight. `Buffer.tableLookupForward` composes exactly this sequence of float32 +primitives, so a native result that agrees with this spec bit for bit inherits the scalar facts +of `Float32Contract`. +-/ +def tableLookupSpec {size k : Nat} (table : FlatBuffer size) + (lower upper : Fin k → Fin size) (weight : FlatBuffer k) : FlatBuffer k := + fun i => + let lowerSample := table (lower i) + ExecFloat.add lowerSample + (ExecFloat.mul (ExecFloat.sub (table (upper i)) lowerSample) (weight i)) + /-! ## Batched row-major matrix multiplication -/ /-- Linear index for either BMM input, viewed as row-major `(batch, rows, cols)` storage. -/ diff --git a/NN/Runtime/Autograd/Engine/LibTorch/Kernels.lean b/NN/Runtime/Autograd/Engine/LibTorch/Kernels.lean index 5df9b66e..c6a38a3d 100644 --- a/NN/Runtime/Autograd/Engine/LibTorch/Kernels.lean +++ b/NN/Runtime/Autograd/Engine/LibTorch/Kernels.lean @@ -363,6 +363,26 @@ Indices that fit in `UInt32` but are out of bounds are ignored; repeated indices opaque scatterAddRows (mat values : @& Buffer) (rows cols : UInt32) (indices : @& Array Nat) (k : UInt32) : Buffer +/-- +Gather `table[positions[i]]` for `count` device-resident positions. + +`table` has `size > 0` entries. Positions are float32 values computed on device: a NaN reads as +`0`, every position is clamped to `[0, size - 1]` and then truncated, so a non-integral position +floors and no selection leaves the table. The result has `count` entries. +-/ +@[never_extract, extern "torchlean_cuda_buffer_gather_at"] +opaque gatherAt (table : @& Buffer) (size : UInt32) (positions : @& Buffer) (count : UInt32) : + Buffer + +/-- +Scatter-add `count` values into a fresh zero buffer of `size > 0` entries at device-resident +positions, addressed exactly as in `gatherAt`: the VJP of `gatherAt` with respect to its table. +Repeated positions accumulate; the native accumulation order is unspecified unless deterministic +algorithms are selected. +-/ +@[never_extract, extern "torchlean_cuda_buffer_scatter_add_at"] +opaque scatterAddAt (size : UInt32) (positions values : @& Buffer) (count : UInt32) : Buffer + end Buffer end LibTorch diff --git a/NN/Runtime/Autograd/Engine/LibTorch/Ops.lean b/NN/Runtime/Autograd/Engine/LibTorch/Ops.lean index 8715e73a..6ff52bbf 100644 --- a/NN/Runtime/Autograd/Engine/LibTorch/Ops.lean +++ b/NN/Runtime/Autograd/Engine/LibTorch/Ops.lean @@ -15,6 +15,7 @@ public import NN.Runtime.Autograd.Engine.LibTorch.Ops.Linear public import NN.Runtime.Autograd.Engine.LibTorch.Ops.NormSoftmax public import NN.Runtime.Autograd.Engine.LibTorch.Ops.SelectiveScan public import NN.Runtime.Autograd.Engine.LibTorch.Ops.Shape +public import NN.Runtime.Autograd.Engine.LibTorch.Ops.TableLookup /-! CUDA operation dispatch surface. diff --git a/NN/Runtime/Autograd/Engine/LibTorch/Ops/Elementwise.lean b/NN/Runtime/Autograd/Engine/LibTorch/Ops/Elementwise.lean index 1e8f2784..1980ea3d 100644 --- a/NN/Runtime/Autograd/Engine/LibTorch/Ops/Elementwise.lean +++ b/NN/Runtime/Autograd/Engine/LibTorch/Ops/Elementwise.lean @@ -68,6 +68,12 @@ returned buffers are owned by the tape/gradient accumulator; workspace buffers a (forward := Buffer.abs) (backward := fun x dLdy => Buffer.absBwd x dLdy) +/-- Pointwise floor node. The function is piecewise constant, so its backward is zero. -/ +@[inline] def floor {s : Shape} (t : Tape) (xId : Nat) : Result (Tape × Nat) := + unary (t := t) "floor" xId s s + (forward := Buffer.floor) + (backward := fun _x dLdy => Buffer.zeros (Buffer.size dLdy)) + /-- Pointwise square-root node using the CUDA buffer derivative convention. -/ @[inline] def sqrt {s : Shape} (t : Tape) (xId : Nat) : Result (Tape × Nat) := unary (t := t) "sqrt" xId s s diff --git a/NN/Runtime/Autograd/Engine/LibTorch/Ops/Indexing.lean b/NN/Runtime/Autograd/Engine/LibTorch/Ops/Indexing.lean index 70bd052b..1140b15a 100644 --- a/NN/Runtime/Autograd/Engine/LibTorch/Ops/Indexing.lean +++ b/NN/Runtime/Autograd/Engine/LibTorch/Ops/Indexing.lean @@ -226,6 +226,38 @@ end Indexing (sourceId, { s := sourceShape, buf := dSource })] } pure (t.addNode node) +/-! +## Device-resident positions +-/ + +/-- +Gather entries of a rank-one table at positions held in another tape value. + +The positions are float32 values computed on device and addressed as `Buffer.gatherAt` documents. +The result takes the positions' shape. The table receives the scatter-add VJP; the positions are +selectors rather than differentiable inputs, so no gradient flows to them. +-/ +@[inline] def gatherAt {size : Nat} {s : Shape} (t : Tape) (tableId positionsId : Nat) : + Result (Tape × Nat) := do + if size = 0 then + throw "autograd: gather_at: the table is empty" + let tableShape : Shape := .dim size .scalar + let table ← requireValue (t := t) tableId tableShape + let positions ← requireValue (t := t) positionsId s + let size32 ← AnyBuffer.natToU32Checked size + let count32 ← AnyBuffer.numelU32 s + let y := Buffer.gatherAt table size32 positions count32 + let node : Node := + { name := some "gather_at" + value := { s := s, buf := y } + requiresGrad := (t.getNode? tableId).any (·.requiresGrad) + parents := #[tableId, positionsId] + backward := fun dLdyAny => do + let dLdy ← requireGrad dLdyAny s + let dTable := Buffer.scatterAddAt size32 positions dLdy.buf count32 + pure #[(tableId, { s := tableShape, buf := dTable })] } + pure (t.addNode node) + end Tape end LibTorch diff --git a/NN/Runtime/Autograd/Engine/LibTorch/Ops/TableLookup.lean b/NN/Runtime/Autograd/Engine/LibTorch/Ops/TableLookup.lean new file mode 100644 index 00000000..5042d8aa --- /dev/null +++ b/NN/Runtime/Autograd/Engine/LibTorch/Ops/TableLookup.lean @@ -0,0 +1,178 @@ +/- +Copyright (c) 2026 TorchLean +Released under MIT license as described in the file LICENSE. +Authors: TorchLean Team +-/ + +module + +public import NN.Runtime.Autograd.Engine.LibTorch.Ops.Core + +/-! +# CUDA Tape Operations: Piecewise-Linear Table Lookup + +A function tabulated at `width` uniformly spaced grid points is evaluated at a grid coordinate +`u` by linear interpolation between the two samples around `u`, after clamping `u` to +`[0, width - 1]`. This is the uniform-grid analogue of `np.interp`, which is how calibration +curves, inverted sensor-response tables, and activation lookups are evaluated. A layered table +stores one such function per layer and selects each element's layer from a second buffer. + +With `i = ⌊u⌋` and `w = u - i`, the value is `v₀ + (v₁ - v₀) · w` for the samples `v₀ = table[i]` +and `v₁ = table[min (i + 1) (width - 1)]`. The forward composes numerical primitives that each +round once in float32, so the result is `KernelSpec.tableLookupSpec` bit for bit. The backward is +TorchLean's own: the coordinate receives the slope of its segment, zero where the clamp is active +as for `clamp`, and the table receives the two interpolation weights scattered to the samples +they multiplied. +-/ + +@[expose] public section + +namespace Runtime +namespace Autograd +namespace LibTorch + +open Spec TorchLean + +namespace Buffer + +/-- The samples each element interpolates between, computed once and reused by the backward. -/ +structure TableSegments where + /-- Flat position of the lower sample. -/ + lower : Buffer + /-- Flat position of the upper sample; the lower one again on the last node of a layer. -/ + upper : Buffer + /-- Interpolation weight toward the upper sample, in `[0, 1)`. -/ + weight : Buffer + /-- Intermediates that the forward value releases once it has consumed them. -/ + workspace : Array Buffer := #[] + +/-- +Locate the segment of each of `count` coordinates inside a table of `layers × width` samples. + +Coordinates clamp to `[0, width - 1]` and the optional layer indices to `[0, layers - 1]` with the +backend's `clamp`, whose NaN selects the lower bound. Every operation here is exact in float32 +while `layers * width ≤ 2^24`; the tape operation enforces that bound. +-/ +@[no_expose] def tableSegments (layers width : UInt32) (coords : Buffer) (layer : Option Buffer) + (count : UInt32) : TableSegments := + let last : Float := Float.ofNat (width.toNat - 1) + let coordinate := clamp coords 0.0 last + let column := floor coordinate + let weight := releaseThen coordinate (sub coordinate column) + let ones := full count 1.0 + let next := releaseThen ones (add column ones) + let nextColumn := releaseThen next (clamp next 0.0 last) + match layer with + | none => { lower := column, upper := nextColumn, weight := weight } + | some layerIdx => + let lastLayer : Float := Float.ofNat (layers.toNat - 1) + let clampedLayer := clamp layerIdx 0.0 lastLayer + let integralLayer := releaseThen clampedLayer (floor clampedLayer) + let offset := releaseThen integralLayer (scale integralLayer (Float.ofNat width.toNat)) + let lower := add offset column + let upper := add offset nextColumn + { lower := lower, upper := upper, weight := weight + workspace := #[column, nextColumn, offset] } + +/-- +Interpolate the table at located segments: `v₀ + (v₁ - v₀) · w`, one rounding per operation. + +The segments' workspace is released through the result, so the forward value is the last use of +those intermediates. +-/ +@[no_expose] def tableLookupForward (table : Buffer) (size : UInt32) (seg : TableSegments) + (count : UInt32) : Buffer := + let lowerSample := gatherAt table size seg.lower count + let upperSample := gatherAt table size seg.upper count + let slope := sub upperSample lowerSample + let scaled := mul slope seg.weight + let value := add lowerSample scaled + releaseManyThen (#[lowerSample, upperSample, slope, scaled] ++ seg.workspace) value + +/-- +VJP of `tableLookupForward`, returning the coordinate and table cotangents. + +The coordinate cotangent is the segment slope times the upstream cotangent, zeroed where the +coordinate clamp was active (the open-interval rule of `clampBwd`). The table cotangent scatters +`(1 - w)` and `w` times the upstream cotangent to the lower and upper samples. +-/ +@[no_expose] def tableLookupBackward (table : Buffer) (size : UInt32) (seg : TableSegments) + (coords : Buffer) (width : UInt32) (dLdy : Buffer) (count : UInt32) : Buffer × Buffer := + let last : Float := Float.ofNat (width.toNat - 1) + let lowerSample := gatherAt table size seg.lower count + let upperSample := gatherAt table size seg.upper count + let slope := sub upperSample lowerSample + let interior := mul dLdy slope + let dCoords := releaseManyThen #[lowerSample, upperSample, slope, interior] + (clampBwd coords interior 0.0 last) + let ones := full count 1.0 + let lowerWeight := sub ones seg.weight + let dLower := mul dLdy lowerWeight + let dUpper := mul dLdy seg.weight + let lowerPart := scatterAddAt size seg.lower dLower count + let upperPart := scatterAddAt size seg.upper dUpper count + let dTable := releaseManyThen #[ones, lowerWeight, dLower, dUpper, lowerPart, upperPart] + (add lowerPart upperPart) + (dCoords, dTable) + +end Buffer + +namespace Tape + +/-- Shared implementation of both lookups; `layerId` supplies per-element layer indices. -/ +@[no_expose] def tableLookupCore (layers width : Nat) {s : Shape} (t : Tape) + (tableId coordsId : Nat) (layerId : Option Nat) : Result (Tape × Nat) := do + if width = 0 then + throw "autograd: table_lookup: the table has no samples" + if layers = 0 then + throw "autograd: table_lookup: the table has no layers" + if layers * width > 2 ^ 24 then + throw "autograd: table_lookup: float32 positions address at most 2^24 samples exactly" + let tableShape : Shape := match layerId with + | none => .dim width .scalar + | some _ => .dim layers (.dim width .scalar) + let table ← requireValue (t := t) tableId tableShape + let coords ← requireValue (t := t) coordsId s + let layer ← match layerId with + | none => pure none + | some id => some <$> requireValue (t := t) id s + let layers32 ← AnyBuffer.natToU32Checked layers + let width32 ← AnyBuffer.natToU32Checked width + let size32 ← AnyBuffer.natToU32Checked (layers * width) + let count32 ← AnyBuffer.numelU32 s + let seg := Buffer.tableSegments layers32 width32 coords layer count32 + let y := Buffer.tableLookupForward table size32 seg count32 + let node : Node := + { name := some (if layerId.isSome then "layered_table_lookup" else "table_lookup") + value := { s := s, buf := y } + requiresGrad := (t.getNode? tableId).any (·.requiresGrad) || + (t.getNode? coordsId).any (·.requiresGrad) + parents := #[tableId, coordsId] ++ (match layerId with | none => #[] | some id => #[id]) + cleanup := #[seg.lower, seg.upper, seg.weight] + backward := fun dLdyAny => do + let dLdy ← requireGrad dLdyAny s + let (dCoords, dTable) := + Buffer.tableLookupBackward table size32 seg coords width32 dLdy.buf count32 + pure #[ + (tableId, { s := tableShape, buf := dTable }), + (coordsId, { s := s, buf := dCoords })] } + pure (t.addNode node) + +/-- Piecewise-linear lookup of a rank-one table of `width` samples at coordinates of shape `s`. -/ +@[inline] def tableLookup {width : Nat} {s : Shape} (t : Tape) (tableId coordsId : Nat) : + Result (Tape × Nat) := + tableLookupCore 1 width (s := s) t tableId coordsId none + +/-- +Piecewise-linear lookup of a `layers × width` table, selecting each element's layer from the +integral values of the tape value `layerId`, which has shape `s` and receives no gradient. +-/ +@[inline] def layeredTableLookup {layers width : Nat} {s : Shape} (t : Tape) + (tableId coordsId layerId : Nat) : Result (Tape × Nat) := + tableLookupCore layers width (s := s) t tableId coordsId (some layerId) + +end Tape + +end LibTorch +end Autograd +end Runtime diff --git a/NN/Tests/Runtime/Cuda/Suite.lean b/NN/Tests/Runtime/Cuda/Suite.lean index b9cb35c5..21968b34 100644 --- a/NN/Tests/Runtime/Cuda/Suite.lean +++ b/NN/Tests/Runtime/Cuda/Suite.lean @@ -10,6 +10,7 @@ public import NN.Tests.Runtime.ParameterAliases public import NN.Tests.Runtime.Cuda.Softmax public import NN.Tests.Runtime.Cuda.Elementwise public import NN.Tests.Runtime.Cuda.Attention +public import NN.Tests.Runtime.Cuda.TableLookup public import NN.Tests.Runtime.Cuda.SelectiveScan public import NN.Tests.Runtime.Cuda.Fft public import NN.Tests.Runtime.Cuda.Stress @@ -37,6 +38,7 @@ def run : IO Unit := do Softmax.run Elementwise.run Attention.run + TableLookup.run SelectiveScan.run Fft.run Stress.run diff --git a/NN/Tests/Runtime/Cuda/TableLookup.lean b/NN/Tests/Runtime/Cuda/TableLookup.lean new file mode 100644 index 00000000..05d0b1db --- /dev/null +++ b/NN/Tests/Runtime/Cuda/TableLookup.lean @@ -0,0 +1,352 @@ +/- +Copyright (c) 2026 TorchLean +Released under MIT license as described in the file LICENSE. +Authors: TorchLean Team +-/ + +module + +public import NN.Runtime.Autograd.Engine.LibTorch.Ops +public import NN.Tensor +public import NN.Tests.Runtime.Cuda.Utils + +/-! +# CUDA Kernel Coverage: Floor, Device-Position Gather, Table Lookup + +- `floor` agrees with `Float32.floor` bit for bit, and its tape node has a zero gradient; +- `gatherAt` agrees with the CPU tape's `indexSelect` for in-range positions, forward and table + gradient, and follows its documented addressing for out-of-range and NaN positions; +- `tableLookup` and `layeredTableLookup` reproduce an executable `Float32` reference bit for bit + at interior, boundary, out-of-range, integral, and NaN coordinates, and their forward and both + gradients agree with the same function recorded on the CPU tape from existing nodes; +- `scatterAddAt` accumulating many updates into one entry, a worst case for the order of + floating-point additions, returns the same bits on a repeated call under LibTorch's + deterministic mode. This is a runtime regression on the current build and device, not a proof + of an accumulation order. +-/ + +@[expose] public section + +namespace Tests +namespace Cuda +namespace TableLookup + +open Spec TorchLean +open TorchLean TorchLean.Tensor +open Runtime.Autograd +open Runtime.Autograd.LibTorch (Buffer) + +/-- Bitwise comparison of two float arrays after narrowing both to float32. -/ +def assertBits32 (msg : String) (got expected : FloatArray) : IO Unit := do + unless got.size == expected.size do + throw <| IO.userError s!"{msg}: size {got.size}, expected {expected.size}" + for i in [:got.size] do + let x := (got.get! i).toFloat32 + let y := (expected.get! i).toFloat32 + unless x.toBits == y.toBits do + throw <| IO.userError + s!"{msg}[{i}]: got {x} (bits {x.toBits}), expected {y} (bits {y.toBits})" + +/-- Read a CUDA tape value as a raw float array. -/ +def rawValue (t : LibTorch.Tape) (id : Nat) (s : Shape) : IO FloatArray := do + let b ← IO.ofExcept (LibTorch.Tape.requireValue (t := t) id s) + pure (Buffer.toFloatArray b) + +def nan : Float := Float.ofBits 0x7ff8000000000000 + +def runFloor : IO Unit := do + IO.println "== floor ==" + let inputs : Array Float := #[-2.5, -1.0, -0.5, -0.0, 0.0, 0.5, 1.0, 1.5, 2.9999, + 16777217.0, -16777217.0, 1.0e8, -1.0e8, 0.1, -0.1] + let expected := FloatArray.mk <| inputs.map fun x => x.toFloat32.floor.toFloat + let input := Buffer.ofFloatArray (FloatArray.mk inputs) + assertBits32 "floor buffer" (Buffer.toFloatArray (Buffer.floor input)) expected + let s : Shape := [inputs.size] + let t0 : LibTorch.Tape := LibTorch.Tape.empty + let (t1, xId) := LibTorch.Tape.leaf (t := t0) { s := s, buf := input } (name := some "x") + let (t2, yId) ← IO.ofExcept (LibTorch.Tape.floor (s := s) (t := t1) xId) + assertBits32 "floor tape value" (← rawValue t2 yId s) expected + let seed : LibTorch.AnyBuffer := { s := s, buf := Buffer.full inputs.size.toUInt32 1.0 } + let grads ← IO.ofExcept (LibTorch.Tape.backwardDenseAll (t := t2) yId seed) + let some dx := grads[xId]? + | throw <| IO.userError "floor: missing gradient" + assertBits32 "floor gradient" (Buffer.toFloatArray dx.buf) + (FloatArray.mk (Array.replicate inputs.size 0.0)) + +def runGatherAt : IO Unit := do + IO.println "== gatherAt vs CPU indexSelect ==" + let n : Nat := 5 + let k : Nat := 4 + let sT : Shape := [n] + let sP : Shape := [k] + let tableVals : Array Float := #[0.125, -0.25, 0.5, 1.0, -2.0] + let table : Tensor Float sT := (Tensor.from tableVals).reshape [n] (by dsimp; decide) + let positions : Fin k → Fin n := + ![⟨2, by decide⟩, ⟨0, by decide⟩, ⟨2, by decide⟩, ⟨4, by decide⟩] + let idx : Tensor (Fin n) [k] := Tensor.ofFn positions + let positionsF : Tensor Float sP := Tensor.ofFn fun i => Float.ofNat (positions i).val + let seed : Tensor Float sP := (Tensor.from #[1.0, 0.5, -0.25, 2.0]).reshape [k] (by dsimp; decide) + + -- CPU + let t0 : Tape Float := Tape.empty + let (t1, tId) := Tape.leaf (t := t0) table (name := some "table") + let (t2, yId) ← IO.ofExcept + (Tape.indexSelect (α := Float) (s := sT) (t := t1) tId 0 k idx) + let yCpu ← Utils.cpuValue (s := sP) t2 yId + let gradsCpu ← IO.ofExcept + (Tape.backwardDenseAll (α := Float) (t := t2) yId (Spec.SomeTensor.ofTensor seed)) + let dTCpu ← Utils.cpuGrad (s := sT) gradsCpu tId + + -- CUDA + let c0 : LibTorch.Tape := LibTorch.Tape.empty + let (c1, tIdc) := LibTorch.Tape.leaf (t := c0) (Utils.tensorToAnyBuffer table) + (name := some "table") + let (c2, pIdc) := LibTorch.Tape.leaf (t := c1) (Utils.tensorToAnyBuffer positionsF) + (name := some "positions") (requiresGrad := false) + let (c3, yIdc) ← IO.ofExcept + (LibTorch.Tape.gatherAt (size := n) (s := sP) (t := c2) tIdc pIdc) + let yCuda ← Utils.cudaValue (s := sP) c3 yIdc + let gradsCuda ← IO.ofExcept + (LibTorch.Tape.backwardDenseAll (t := c3) yIdc (Utils.tensorToAnyBuffer seed)) + let dTCuda ← Utils.cudaGrad (s := sT) gradsCuda tIdc + + Utils.assertTensorApprox (s := sP) "gather_at forward" yCuda yCpu (tol := 0.0) + Utils.assertTensorApprox (s := sT) "gather_at table gradient" dTCuda dTCpu (tol := 0.0) + + -- Addressing: negative, past the end, non-integral, NaN, and a negative fraction. + let tableBuf := Buffer.ofFloatArray (FloatArray.mk tableVals) + let odd := Buffer.ofFloatArray (FloatArray.mk #[-3.0, 7.0, 2.75, nan, -0.5]) + assertBits32 "gather_at addressing" (Buffer.toFloatArray (Buffer.gatherAt tableBuf 5 odd 5)) + (FloatArray.mk #[0.125, -2.0, 0.5, 0.125, 0.125]) + +/-- The backend's clamp, `fmin (fmax x lo) hi`, which selects `lo` for a NaN input. -/ +def clampRef (x lo hi : Float32) : Float32 := + let m := if x.isNaN then lo else if x < lo then lo else x + if m > hi then hi else m + +/-- +Executable reference for one lookup in Lean core `Float32` (IEEE binary32), statement for +statement the sequence `Buffer.tableSegments` and `Buffer.tableLookupForward` run on device. +-/ +def refLookup (table : Array Float32) (layers width : Nat) (u layer : Float32) : Float32 := + let last := Float32.ofNat (width - 1) + let coordinate := clampRef u 0.0 last + let column := coordinate.floor + let weight := coordinate - column + let nextColumn := clampRef (column + 1.0) 0.0 last + let layerIdx := (clampRef layer 0.0 (Float32.ofNat (layers - 1))).floor + let offset := layerIdx * Float32.ofNat width + let lower := (offset + column).toFloat.toUInt64.toNat + let upper := (offset + nextColumn).toFloat.toUInt64.toNat + let lowerSample := table[lower]! + let upperSample := table[upper]! + lowerSample + (upperSample - lowerSample) * weight + +/-- Smooth, layer-dependent samples in layer-major order. -/ +def sampleData (width layers : Nat) : FloatArray := Id.run do + let mut a := FloatArray.emptyWithCapacity (width * layers) + for l in [0:layers] do + for i in [0:width] do + let x := Float.ofNat i + let y := Float.ofNat l + a := a.push (0.17 * x - 0.031 * x * x + 0.4 * y + 0.05) + return a + +def widthN : Nat := 7 +def layersN : Nat := 3 + +/-- Interior, boundary, out-of-range, integral, signed-zero, and NaN grid coordinates. -/ +def probeCoords : Array Float := + #[0.0, 0.25, 1.0, 2.5, 3.75, 5.999, 6.0, -1.5, 9.75, 4.125, 2.875, 0.001, nan, -0.0, 5.5] + +/-- Layer selectors, including one past the end and one non-integral. -/ +def probeLayers : Array Float := + #[0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 1, 5.0, 1.5] + +def runLookupBitExact : IO Unit := do + IO.println "== tableLookup bit-exact against the Float32 reference ==" + let count := probeCoords.size + let sX : Shape := [count] + let data := sampleData widthN layersN + let table32 := (Array.range data.size).map fun i => (data.get! i).toFloat32 + let coordsBuf := Buffer.ofFloatArray (FloatArray.mk probeCoords) + let layersBuf := Buffer.ofFloatArray (FloatArray.mk probeLayers) + + -- Layered table. + let expectedLayered := FloatArray.mk <| (Array.range count).map fun i => + (refLookup table32 layersN widthN (probeCoords[i]!).toFloat32 + (probeLayers[i]!).toFloat32).toFloat + let c0 : LibTorch.Tape := LibTorch.Tape.empty + let (c1, tId) := LibTorch.Tape.leaf (t := c0) + { s := [layersN, widthN], buf := Buffer.ofFloatArray data } (name := some "table") + let (c2, uId) := LibTorch.Tape.leaf (t := c1) { s := sX, buf := coordsBuf } (name := some "u") + let (c3, lId) := LibTorch.Tape.leaf (t := c2) { s := sX, buf := layersBuf } + (name := some "layer") (requiresGrad := false) + let (c4, yId) ← IO.ofExcept + (LibTorch.Tape.layeredTableLookup (layers := layersN) (width := widthN) (s := sX) (t := c3) + tId uId lId) + assertBits32 "layered_table_lookup" (← rawValue c4 yId sX) expectedLayered + + -- Rank-one table: the first layer alone, every probe addressing layer 0. + let layer0 := FloatArray.mk <| (Array.range widthN).map fun i => data.get! i + let table0 := (Array.range widthN).map fun i => (layer0.get! i).toFloat32 + let expectedFlat := FloatArray.mk <| (Array.range count).map fun i => + (refLookup table0 1 widthN (probeCoords[i]!).toFloat32 0.0).toFloat + let (d1, t0Id) := LibTorch.Tape.leaf (t := c0) + { s := [widthN], buf := Buffer.ofFloatArray layer0 } (name := some "table") + let (d2, u0Id) := LibTorch.Tape.leaf (t := d1) { s := sX, buf := coordsBuf } (name := some "u") + let (d3, y0Id) ← IO.ofExcept + (LibTorch.Tape.tableLookup (width := widthN) (s := sX) (t := d2) t0Id u0Id) + assertBits32 "table_lookup" (← rawValue d3 y0Id sX) expectedFlat + +def runLookupVsCpu : IO Unit := do + IO.println "== tableLookup vs the same function on the CPU tape ==" + let width : Nat := 6 + let n : Nat := 8 + let sT : Shape := [width] + let sX : Shape := [n] + -- Dyadic values keep every intermediate exact in both float32 and Float. + let tableVals : Array Float := #[0.5, -1.0, 2.0, 1.5, -0.25, 3.0] + let coordVals : Array Float := #[0.0, 0.25, 1.5, 2.0, 3.75, 5.0, -2.0, 7.5] + let seedVals : Array Float := #[1.0, -0.5, 0.25, 2.0, -1.0, 0.5, 1.5, -0.75] + let table : Tensor Float sT := (Tensor.from tableVals).reshape [width] (by dsimp; decide) + let coords : Tensor Float sX := (Tensor.from coordVals).reshape [n] (by dsimp; decide) + let seed : Tensor Float sX := (Tensor.from seedVals).reshape [n] (by dsimp; decide) + let last : Float := Float.ofNat (width - 1) + -- Host-side segment selection, the recipe the device follows. + let segment : Fin n → Nat × Nat := fun i => + let u := coordVals[i.val]! + let c := if u < 0.0 then 0.0 else if u > last then last else u + let lower := c.floor.toUInt64.toNat + (lower, min (lower + 1) (width - 1)) + let lowerIdx : Tensor (Fin width) [n] := Tensor.ofFn fun i => + ⟨(segment i).1 % width, Nat.mod_lt _ (by decide)⟩ + let upperIdx : Tensor (Fin width) [n] := Tensor.ofFn fun i => + ⟨(segment i).2 % width, Nat.mod_lt _ (by decide)⟩ + let column : Tensor Float sX := Tensor.ofFn fun i => Float.ofNat (segment i).1 + + -- CPU: clamp, subtract the constant column, gather both samples, interpolate. + let t0 : Tape Float := Tape.empty + let (t1, tId) := Tape.leaf (t := t0) table (name := some "table") + let (t2, uId) := Tape.leaf (t := t1) coords (name := some "coords") + let (t3, cId) := Tape.leaf (t := t2) column (name := some "column") (requiresGrad := false) + let (t4, ucId) ← IO.ofExcept (Tape.clamp (α := Float) (s := sX) (t := t3) uId 0.0 last) + let (t5, wId) ← IO.ofExcept (Tape.sub (α := Float) (s := sX) (t := t4) ucId cId) + let (t6, v0Id) ← IO.ofExcept + (Tape.indexSelect (α := Float) (s := sT) (t := t5) tId 0 n lowerIdx) + let (t7, v1Id) ← IO.ofExcept + (Tape.indexSelect (α := Float) (s := sT) (t := t6) tId 0 n upperIdx) + let (t8, slopeId) ← IO.ofExcept (Tape.sub (α := Float) (s := sX) (t := t7) v1Id v0Id) + let (t9, scaledId) ← IO.ofExcept (Tape.mul (α := Float) (s := sX) (t := t8) slopeId wId) + let (t10, yId) ← IO.ofExcept (Tape.add (α := Float) (s := sX) (t := t9) v0Id scaledId) + let yCpu ← Utils.cpuValue (s := sX) t10 yId + let gradsCpu ← IO.ofExcept + (Tape.backwardDenseAll (α := Float) (t := t10) yId (Spec.SomeTensor.ofTensor seed)) + let dTCpu ← Utils.cpuGrad (s := sT) gradsCpu tId + let dUCpu ← Utils.cpuGrad (s := sX) gradsCpu uId + + -- CUDA + let c0 : LibTorch.Tape := LibTorch.Tape.empty + let (c1, tIdc) := LibTorch.Tape.leaf (t := c0) (Utils.tensorToAnyBuffer table) + (name := some "table") + let (c2, uIdc) := LibTorch.Tape.leaf (t := c1) (Utils.tensorToAnyBuffer coords) + (name := some "coords") + let (c3, yIdc) ← IO.ofExcept + (LibTorch.Tape.tableLookup (width := width) (s := sX) (t := c2) tIdc uIdc) + let yCuda ← Utils.cudaValue (s := sX) c3 yIdc + let gradsCuda ← IO.ofExcept + (LibTorch.Tape.backwardDenseAll (t := c3) yIdc (Utils.tensorToAnyBuffer seed)) + let dTCuda ← Utils.cudaGrad (s := sT) gradsCuda tIdc + let dUCuda ← Utils.cudaGrad (s := sX) gradsCuda uIdc + + Utils.assertTensorApprox (s := sX) "table_lookup forward" yCuda yCpu (tol := 1e-6) + Utils.assertTensorApprox (s := sT) "table_lookup table gradient" dTCuda dTCpu (tol := 1e-6) + Utils.assertTensorApprox (s := sX) "table_lookup coordinate gradient" dUCuda dUCpu + (tol := 1e-6) + +def runEdges : IO Unit := do + IO.println "== tableLookup edges (width 1, empty coordinates, rejected shapes) ==" + -- A single sample: every coordinate reads it, and no coordinate has a gradient. + let one : Shape := [1] + let three : Shape := [3] + let c0 : LibTorch.Tape := LibTorch.Tape.empty + let (c1, tId) := LibTorch.Tape.leaf (t := c0) + { s := one, buf := Buffer.ofFloatArray (FloatArray.mk #[0.75]) } (name := some "table") + let (c2, uId) := LibTorch.Tape.leaf (t := c1) + { s := three, buf := Buffer.ofFloatArray (FloatArray.mk #[0.0, 3.5, -2.0]) } + (name := some "u") + let (c3, yId) ← IO.ofExcept + (LibTorch.Tape.tableLookup (width := 1) (s := three) (t := c2) tId uId) + assertBits32 "width 1 value" (← rawValue c3 yId three) (FloatArray.mk #[0.75, 0.75, 0.75]) + let seed : LibTorch.AnyBuffer := + { s := three, buf := Buffer.ofFloatArray (FloatArray.mk #[1.0, 2.0, 4.0]) } + let grads ← IO.ofExcept (LibTorch.Tape.backwardDenseAll (t := c3) yId seed) + let some dT := grads[tId]? + | throw <| IO.userError "width 1: missing table gradient" + let some dU := grads[uId]? + | throw <| IO.userError "width 1: missing coordinate gradient" + assertBits32 "width 1 table gradient" (Buffer.toFloatArray dT.buf) (FloatArray.mk #[7.0]) + assertBits32 "width 1 coordinate gradient" (Buffer.toFloatArray dU.buf) + (FloatArray.mk #[0.0, 0.0, 0.0]) + + -- No coordinates: an empty result and empty gradients, with nothing read from the table. + let none : Shape := [0] + let (e1, eId) := LibTorch.Tape.leaf (t := c1) + { s := none, buf := Buffer.ofFloatArray (FloatArray.mk #[]) } (name := some "u") + let (e2, eyId) ← IO.ofExcept + (LibTorch.Tape.tableLookup (width := 1) (s := none) (t := e1) tId eId) + unless (← rawValue e2 eyId none).size == 0 do + throw <| IO.userError "empty coordinates: expected an empty value" + let emptySeed : LibTorch.AnyBuffer := { s := none, buf := Buffer.zeros 0 } + let emptyGrads ← IO.ofExcept (LibTorch.Tape.backwardDenseAll (t := e2) eyId emptySeed) + let some edT := emptyGrads[tId]? + | throw <| IO.userError "empty coordinates: missing table gradient" + assertBits32 "empty coordinates table gradient" (Buffer.toFloatArray edT.buf) + (FloatArray.mk #[0.0]) + + -- Rejected before any buffer is read: no samples, and more samples than float32 addresses. + match LibTorch.Tape.tableLookup (width := 0) (s := three) (t := c2) tId uId with + | .ok _ => throw <| IO.userError "width 0: expected an error" + | .error msg => + unless (msg.splitOn "no samples").length > 1 do + throw <| IO.userError s!"width 0: unexpected error: {msg}" + match LibTorch.Tape.layeredTableLookup (layers := 4097) (width := 4096) (s := three) + (t := c2) tId uId uId with + | .ok _ => throw <| IO.userError "2^24 bound: expected an error" + | .error msg => + unless (msg.splitOn "2^24").length > 1 do + throw <| IO.userError s!"2^24 bound: unexpected error: {msg}" + +/-- Repeatability of a device-position scatter-add under the deterministic control. -/ +def runScatterAddAtTwice : IO Unit := do + IO.println "== deterministic scatter_add_at: exact repeatability ==" + let deterministic ← LibTorch.getDeterministic + let benchmark ← LibTorch.getCuDNNBenchmark + try + LibTorch.setDeterministic true + let k : UInt32 := 4096 + let one : UInt32 := 1 + let big := Buffer.full one 1.0e8 + let ones := Buffer.full (k - one) 1.0 + let values := Buffer.concatBuffers big ones one (k - one) + let positions := Buffer.zeros k + let y1 := Buffer.scatterAddAt one positions values k + let y2 := Buffer.scatterAddAt one positions values k + Utils.assertFloatArrayEq "scatterAddAt deterministic run1 vs run2" + (Buffer.toFloatArray y1) (Buffer.toFloatArray y2) + finally + LibTorch.setDeterministic deterministic + LibTorch.setCuDNNBenchmark benchmark + +def run : IO Unit := do + IO.println "=== CUDA kernel coverage: floor, gather_at, table lookup ===" + runFloor + runGatherAt + runLookupBitExact + runLookupVsCpu + runEdges + runScatterAddAtTwice + +end TableLookup +end Cuda +end Tests diff --git a/csrc/libtorch/README.md b/csrc/libtorch/README.md index b400310f..c63539b3 100644 --- a/csrc/libtorch/README.md +++ b/csrc/libtorch/README.md @@ -182,6 +182,7 @@ Current CUDA coverage: | `NN/Tests/Runtime/Cuda/Softmax.lean` | `softmax` and `log_softmax`, forward and backward. | | `NN/Tests/Runtime/Cuda/Elementwise.lean` | Scalar elementwise ops, activations, safe logs, products, and `sum`. | | `NN/Tests/Runtime/Cuda/Attention.lean` | Attention values and gradients against the CPU reference, including batches. | +| `NN/Tests/Runtime/Cuda/TableLookup.lean` | `floor`, gather at device-resident positions, and piecewise-linear table lookup: bit-exact forward against the float32 reference, forward and gradients against the CPU tape; `scatterAddAt` repeatability under the deterministic control. | | `NN/Tests/Runtime/Cuda/SelectiveScan.lean` | Diagonal selective-scan buffer primitives used by the Mamba/SSM runtime path. | | `NN/Tests/Runtime/Cuda/Fft.lean` | Packed real FFT, inverse FFT, spectral convolution, and finite-difference gradient checks. | | `NN/Tests/Runtime/Cuda/Stress.lean` | RNG determinism, explicit release, duplicate-parent gradient accumulation, large buffers, reductions, and rectangular matmul. | diff --git a/csrc/libtorch/operations.h b/csrc/libtorch/operations.h index 9b2b36a1..fd16b2c1 100644 --- a/csrc/libtorch/operations.h +++ b/csrc/libtorch/operations.h @@ -2,6 +2,7 @@ // Intentionally included more than once, with the signature macros defined by the caller. TORCHLEAN_UNARY_EXPORT(abs, at::abs(x)) +TORCHLEAN_UNARY_EXPORT(floor, at::floor(x)) TORCHLEAN_UNARY_EXPORT(sqrt, selected_sqrt(x)) TORCHLEAN_UNARY_EXPORT(exp, at::exp(x)) TORCHLEAN_UNARY_EXPORT(sin, at::sin(x)) diff --git a/csrc/libtorch/torchlean.cpp b/csrc/libtorch/torchlean.cpp index 4aaa9d3f..c143f0b2 100644 --- a/csrc/libtorch/torchlean.cpp +++ b/csrc/libtorch/torchlean.cpp @@ -1784,6 +1784,40 @@ extern "C" LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_scatter_add_rows( }); } +namespace { + +// Positions computed on device arrive as float32 values, so no host round trip selects them. +// A NaN position reads as the first entry; every position is then clamped into the table and +// truncated, which floors a non-integral value and keeps each selection in bounds without a +// device-side assertion. +Tensor table_positions(const Tensor& positions, uint32_t size) { + const double last = static_cast(size) - 1.0; + return at::clamp(at::nan_to_num(positions, 0.0, last, 0.0), 0.0, last).to(at::kLong); +} + +} // namespace + +extern "C" LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_gather_at( + b_lean_obj_arg table, uint32_t size, b_lean_obj_arg positions, uint32_t count) { + return invoke([&] { + require(size > 0, "LibTorch kernels: gather_at needs a nonempty table"); + const auto index = table_positions(checked(positions, {count}), size); + return box(checked(table, {size}).index_select(0, index)); + }); +} + +extern "C" LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_scatter_add_at( + uint32_t size, b_lean_obj_arg positions, b_lean_obj_arg values, uint32_t count) { + return invoke([&] { + require(size > 0, "LibTorch kernels: scatter_add_at needs a nonempty table"); + const auto index = table_positions(checked(positions, {count}), size); + const auto source = checked(values, {count}); + auto result = at::zeros({static_cast(size)}, source.options()); + result.index_add_(0, index, source); + return box(result); + }); +} + } // namespace torchlean::operators // Host double-precision matrix multiplication diff --git a/csrc/libtorch/unavailable.c b/csrc/libtorch/unavailable.c index 9867a2a9..d8179457 100644 --- a/csrc/libtorch/unavailable.c +++ b/csrc/libtorch/unavailable.c @@ -212,6 +212,10 @@ UNAVAILABLE(lean_obj_res, torchlean_cuda_buffer_gather_rows, UNAVAILABLE(lean_obj_res, torchlean_cuda_buffer_scatter_add_rows, b_lean_obj_arg MObj, b_lean_obj_arg ValuesObj, uint32_t rows, uint32_t cols, b_lean_obj_arg IdxObj, uint32_t k) +UNAVAILABLE(lean_obj_res, torchlean_cuda_buffer_gather_at, + b_lean_obj_arg TableObj, uint32_t size, b_lean_obj_arg PositionsObj, uint32_t count) +UNAVAILABLE(lean_obj_res, torchlean_cuda_buffer_scatter_add_at, + uint32_t size, b_lean_obj_arg PositionsObj, b_lean_obj_arg ValuesObj, uint32_t count) UNAVAILABLE(uint32_t, torchlean_cuda_buffer_size, b_lean_obj_arg BObj) UNAVAILABLE(uint8_t, torchlean_cuda_buffer_dtype, b_lean_obj_arg BObj) UNAVAILABLE(lean_obj_res, torchlean_cuda_buffer_format, b_lean_obj_arg buffer) From 11d1e3a62386de6240ee85e52a0bf2cea61bb63d Mon Sep 17 00:00:00 2001 From: Nicolas Rouquette Date: Sun, 4 Oct 2026 12:19:48 -0700 Subject: [PATCH 2/2] fix(libtorch): clamp table positions again in the integer domain `table_positions` clamped float32 positions to `[0, size - 1]` and converted the result to int64, but above 2^24 the bound `size - 1` has no float32 representation and rounds up to `size` (16777219 reads as 16777220), so an infinite or oversize position selected one past the end. `gatherAt` and `scatterAddAt` are not bounded by the 2^24 limit that `Tape.tableLookup` enforces for exact positions. The index is now clamped again as int64, where every table index is exact, and the docstrings say so. The regression gathers from and scatter-adds into a table of 16777220 entries at infinity, at 16777220, at the unrepresentable 16777219, at 1e9, and at NaN and negative positions; without the integer clamp it dies in `index_select` with a device-side assert. --- .../Autograd/Engine/LibTorch/Kernels.lean | 5 ++- NN/Tests/Runtime/Cuda/TableLookup.lean | 42 ++++++++++++++++++- csrc/libtorch/README.md | 2 +- csrc/libtorch/torchlean.cpp | 8 +++- 4 files changed, 52 insertions(+), 5 deletions(-) diff --git a/NN/Runtime/Autograd/Engine/LibTorch/Kernels.lean b/NN/Runtime/Autograd/Engine/LibTorch/Kernels.lean index c6a38a3d..cf7131de 100644 --- a/NN/Runtime/Autograd/Engine/LibTorch/Kernels.lean +++ b/NN/Runtime/Autograd/Engine/LibTorch/Kernels.lean @@ -368,7 +368,10 @@ Gather `table[positions[i]]` for `count` device-resident positions. `table` has `size > 0` entries. Positions are float32 values computed on device: a NaN reads as `0`, every position is clamped to `[0, size - 1]` and then truncated, so a non-integral position -floors and no selection leaves the table. The result has `count` entries. +floors and no selection leaves the table. The clamp is applied to the float32 position and again +to the integer index, so it holds for tables above 2^24 entries, where `size - 1` has no float32 +representation; which entries a float32 position can address exactly is a separate question, +bounded by `Tape.tableLookup`. The result has `count` entries. -/ @[never_extract, extern "torchlean_cuda_buffer_gather_at"] opaque gatherAt (table : @& Buffer) (size : UInt32) (positions : @& Buffer) (count : UInt32) : diff --git a/NN/Tests/Runtime/Cuda/TableLookup.lean b/NN/Tests/Runtime/Cuda/TableLookup.lean index 05d0b1db..0b4bb4cf 100644 --- a/NN/Tests/Runtime/Cuda/TableLookup.lean +++ b/NN/Tests/Runtime/Cuda/TableLookup.lean @@ -15,7 +15,8 @@ public import NN.Tests.Runtime.Cuda.Utils - `floor` agrees with `Float32.floor` bit for bit, and its tape node has a zero gradient; - `gatherAt` agrees with the CPU tape's `indexSelect` for in-range positions, forward and table - gradient, and follows its documented addressing for out-of-range and NaN positions; + gradient, and follows its documented addressing for out-of-range and NaN positions, on a table + above 2^24 entries included; - `tableLookup` and `layeredTableLookup` reproduce an executable `Float32` reference bit for bit at interior, boundary, out-of-range, integral, and NaN coordinates, and their forward and both gradients agree with the same function recorded on the CPU tape from existing nodes; @@ -119,6 +120,44 @@ def runGatherAt : IO Unit := do assertBits32 "gather_at addressing" (Buffer.toFloatArray (Buffer.gatherAt tableBuf 5 odd 5)) (FloatArray.mk #[0.125, -2.0, 0.5, 0.125, 0.125]) +/-- +In-bounds addressing above 2^24 entries, where `size - 1` has no float32 representation: a float32 +clamp alone rounds that bound up to `size` (16777219 reads as 16777220), and an infinite or +oversize position then selects one past the end. The native kernels clamp the integer index again. +-/ +def runGatherAtLargeTable : IO Unit := do + IO.println "== gatherAt / scatterAddAt above 2^24 entries ==" + let size : Nat := 16777220 + let size32 := size.toUInt32 + -- A table of zeros, marked at the entries the positions below must select. + let mut data := FloatArray.emptyWithCapacity size + for _ in [0:size] do + data := data.push 0.0 + data := data.set! 0 1.0 + data := data.set! 16777216 2.0 + data := data.set! 16777218 3.0 + data := data.set! (size - 1) 4.0 + let table := Buffer.ofFloatArray data + let inf : Float := 1.0 / 0.0 + -- 16777219.0 narrows to 16777220.0 in float32, one past the end, exactly as the bound does. + let positions := Buffer.ofFloatArray (FloatArray.mk + #[inf, 16777220.0, 16777219.0, 1.0e9, 16777216.0, 16777218.0, -inf, nan, -3.0]) + assertBits32 "gather_at above 2^24" + (Buffer.toFloatArray (Buffer.gatherAt table size32 positions 9)) + (FloatArray.mk #[4.0, 4.0, 4.0, 4.0, 2.0, 3.0, 1.0, 1.0, 1.0]) + -- The scatter lands at the same entries; read it back through positions that address exactly. + let scatterPositions := Buffer.ofFloatArray (FloatArray.mk + #[inf, 16777220.0, 16777216.0, -inf, nan]) + let values := Buffer.ofFloatArray (FloatArray.mk #[1.0, 2.0, 4.0, 8.0, 16.0]) + let scattered := Buffer.scatterAddAt size32 scatterPositions values 5 + unless Buffer.size scattered == size32 do + throw <| IO.userError + s!"scatter_add_at above 2^24: size {Buffer.size scattered}, expected {size}" + let probes := Buffer.ofFloatArray (FloatArray.mk #[inf, 16777216.0, 0.0, 16777218.0]) + assertBits32 "scatter_add_at above 2^24" + (Buffer.toFloatArray (Buffer.gatherAt scattered size32 probes 4)) + (FloatArray.mk #[3.0, 4.0, 24.0, 0.0]) + /-- The backend's clamp, `fmin (fmax x lo) hi`, which selects `lo` for a NaN input. -/ def clampRef (x lo hi : Float32) : Float32 := let m := if x.isNaN then lo else if x < lo then lo else x @@ -342,6 +381,7 @@ def run : IO Unit := do IO.println "=== CUDA kernel coverage: floor, gather_at, table lookup ===" runFloor runGatherAt + runGatherAtLargeTable runLookupBitExact runLookupVsCpu runEdges diff --git a/csrc/libtorch/README.md b/csrc/libtorch/README.md index c63539b3..5ecee85d 100644 --- a/csrc/libtorch/README.md +++ b/csrc/libtorch/README.md @@ -182,7 +182,7 @@ Current CUDA coverage: | `NN/Tests/Runtime/Cuda/Softmax.lean` | `softmax` and `log_softmax`, forward and backward. | | `NN/Tests/Runtime/Cuda/Elementwise.lean` | Scalar elementwise ops, activations, safe logs, products, and `sum`. | | `NN/Tests/Runtime/Cuda/Attention.lean` | Attention values and gradients against the CPU reference, including batches. | -| `NN/Tests/Runtime/Cuda/TableLookup.lean` | `floor`, gather at device-resident positions, and piecewise-linear table lookup: bit-exact forward against the float32 reference, forward and gradients against the CPU tape; `scatterAddAt` repeatability under the deterministic control. | +| `NN/Tests/Runtime/Cuda/TableLookup.lean` | `floor`, gather at device-resident positions (in bounds above 2^24 entries included), and piecewise-linear table lookup: bit-exact forward against the float32 reference, forward and gradients against the CPU tape; `scatterAddAt` repeatability under the deterministic control. | | `NN/Tests/Runtime/Cuda/SelectiveScan.lean` | Diagonal selective-scan buffer primitives used by the Mamba/SSM runtime path. | | `NN/Tests/Runtime/Cuda/Fft.lean` | Packed real FFT, inverse FFT, spectral convolution, and finite-difference gradient checks. | | `NN/Tests/Runtime/Cuda/Stress.lean` | RNG determinism, explicit release, duplicate-parent gradient accumulation, large buffers, reductions, and rectangular matmul. | diff --git a/csrc/libtorch/torchlean.cpp b/csrc/libtorch/torchlean.cpp index c143f0b2..d76da84f 100644 --- a/csrc/libtorch/torchlean.cpp +++ b/csrc/libtorch/torchlean.cpp @@ -1789,10 +1789,14 @@ namespace { // Positions computed on device arrive as float32 values, so no host round trip selects them. // A NaN position reads as the first entry; every position is then clamped into the table and // truncated, which floors a non-integral value and keeps each selection in bounds without a -// device-side assertion. +// device-side assertion. The float32 clamp alone does not keep it in bounds: above 2^24 the +// bound `size - 1` has no float32 representation and rounds up to `size` (16777219 reads as +// 16777220), so an infinite or oversize position would select one past the end. The clamp is +// therefore repeated on the integer index, where every table index is exact. Tensor table_positions(const Tensor& positions, uint32_t size) { const double last = static_cast(size) - 1.0; - return at::clamp(at::nan_to_num(positions, 0.0, last, 0.0), 0.0, last).to(at::kLong); + const auto clamped = at::clamp(at::nan_to_num(positions, 0.0, last, 0.0), 0.0, last); + return at::clamp(clamped.to(at::kLong), int64_t{0}, static_cast(size) - 1); } } // namespace