feat(libtorch): run the bridge on the host when the SDK is CPU-only

NN/Kernel/Runtime.lean

@@ -18,8 +18,9 @@ bounded number of modules per process, keys them by source and context, and laun
 current stream. Input buffers and the module remain alive until the bounds record is checked.
 
 This is an explicit foreign-code boundary, not a Lean theorem about NVRTC or device execution.
-LibTorch continues to own tensor storage and its established numerical operations. CPU-only
-builds resolve the same extern symbols but report an ordinary IO error instead of executing CUDA.
+LibTorch continues to own tensor storage and its established numerical operations. Builds without
+LibTorch resolve the same extern symbols but report an ordinary IO error instead of executing CUDA,
+and so does the bridge on the host (a CPU-only SDK, or `TORCHLEAN_LIBTORCH_DEVICE=cpu`).
 -/
 
 @[expose] public section

NN/Runtime/Autograd/Engine/LibTorch/Buffer.lean

@@ -11,10 +11,11 @@ public import NN.Runtime.Autograd.Engine.LibTorch.Trusted
 public import NN.Kernel.Cuda.Binary
 
 /-!
-# CUDA Buffers
+# LibTorch Buffers
 
-Low-level dtype-preserving operations for the LibTorch CUDA runtime. TorchLean retains its tape and
-selected local VJPs; native calls do not record a LibTorch autograd graph. CUDA builds use
+Low-level dtype-preserving operations for the LibTorch runtime, on a CUDA device or on the host
+(`Runtime.Autograd.LibTorch.deviceKind`). TorchLean retains its tape and selected local VJPs;
+native calls do not record a LibTorch autograd graph. `-K cuda=true` builds use
 `csrc/libtorch/torchlean.cpp`. Builds without LibTorch link `csrc/libtorch/unavailable.c`, which
 reports `.notLinked` and fails every buffer operation.
 -/
@@ -54,38 +55,58 @@ private opaque fail (message : @& String) : Buffer
 
 /-! ### Runtime Availability -/
 
-/-- What implementation sits behind the CUDA FFI symbols in the current process. -/
+/-- What implementation sits behind the buffer FFI symbols in the current process. -/
 inductive RuntimeStatus where
   /-- The default build does not link LibTorch; buffer operations fail. -/
   | notLinked
-  /-- The project was built with LibTorch CUDA and at least one CUDA device is visible. -/
+  /-- The project was built with LibTorch and the bridge selected a device: a CUDA device, or the
+  host (`Runtime.Autograd.LibTorch.deviceKind`). -/
   | nativeAvailable
-  /-- The project was built with LibTorch CUDA, but no usable CUDA device is visible. -/
+  /-- The project was built with LibTorch, but no usable device is selected: a CUDA-enabled SDK
+  with no visible CUDA device, unless `TORCHLEAN_LIBTORCH_DEVICE=cpu` selects the host. -/
   | nativeUnavailable
   deriving DecidableEq, Repr
 
 /-- Raw status word from the C layer; `runtimeStatus` decodes it. -/
 @[never_extract, extern "torchlean_cuda_runtime_status"]
 private opaque runtimeStatusRaw (token : UInt32) : UInt32
 
-/-- Query whether LibTorch is linked and can see a CUDA device. -/
+/-- Query whether LibTorch is linked and has a device to run on. -/
 @[no_expose] def runtimeStatus (token : UInt32 := 0) : RuntimeStatus :=
   match runtimeStatusRaw token with
   | 0 => .notLinked
   | 1 => .nativeAvailable
   | _ => .nativeUnavailable
 
-/-- Require real CUDA execution for a user-selected CUDA session. -/
+/--
+Require the LibTorch bridge with a device selected: a CUDA device, or the host. This is the check
+for a session that needs the bridge whatever its device (the CLI's `cuda` device, a CPU-only SDK
+in CI, `TORCHLEAN_REQUIRE_NATIVE=1` for the test suite); `requireCudaDevice` is the check for a
+session that needs a CUDA device.
+-/
 def requireNativeRuntime : IO Unit :=
   match runtimeStatus with
   | .nativeAvailable => pure ()
   | .notLinked =>
       throw <| IO.userError
-        "CUDA was requested, but this executable was built without LibTorch; \
+        "the LibTorch bridge was requested, but this executable was built without LibTorch; \
           rebuild and run with `-K cuda=true`"
   | .nativeUnavailable =>
       throw <| IO.userError
-        "CUDA was requested and this is a CUDA build, but no usable CUDA device is visible"
+        "the LibTorch bridge was requested and this is a LibTorch build, but no usable CUDA \
+          device is visible; TORCHLEAN_LIBTORCH_DEVICE=cpu selects the host instead"
+
+/--
+Require the LibTorch bridge on a CUDA device. `requireNativeRuntime` is satisfied by the host, so a
+check that means "a GPU ran this" has to ask for the device kind as well: this is what
+`TORCHLEAN_REQUIRE_CUDA=1` means to the test suite.
+-/
+def requireCudaDevice : IO Unit := do
+  requireNativeRuntime
+  if (← deviceKind) == .host then
+    throw <| IO.userError
+      "a CUDA device was required, but the LibTorch bridge runs on the host: the SDK is CPU-only, \
+        or TORCHLEAN_LIBTORCH_DEVICE=cpu selected the host"
 
 /-! ### Allocator Telemetry -/
 

NN/Runtime/Autograd/Engine/LibTorch/Controls.lean

@@ -7,13 +7,19 @@ Authors: TorchLean Team
 module
 
 /-!
-# LibTorch CUDA Runtime Controls
+# LibTorch Runtime Controls
 
 Effectful configuration and readback for TorchLean's LibTorch bridge. Configure these process-wide
 settings before concurrent runtime work. They are runtime requests, not numerical proof evidence:
 IEEE precision and strict determinism do not establish a reference reduction order, correct
 rounding of every operation, or FloatLib bit agreement.
 
+The bridge selects its device once, at its first call: a visible CUDA device when the SDK was
+built with CUDA support, or the host when the SDK is CPU-only. `TORCHLEAN_LIBTORCH_DEVICE=cpu`
+selects the host on any SDK and `TORCHLEAN_LIBTORCH_DEVICE=cuda` insists on a CUDA device; a CUDA
+SDK without a visible device never falls back to the host silently. `deviceKind` reports the
+choice; the CUDA-only controls below say what they do on the host.
+
 TorchLean owns its tape and selected local VJPs. These controls never enable LibTorch autograd.
 Native tensors retain their supported hardware dtypes; this API does not select arbitrary FloatLib
 formats. `version` reports the linked native build rather than a version assumed by Lean.
@@ -41,6 +47,9 @@ private opaque deviceCountRaw (token : UInt32) : UInt32
 @[never_extract, extern "torchlean_libtorch_get_device"]
 private opaque getDeviceRaw (token : UInt32) : UInt32
 
+@[never_extract, extern "torchlean_libtorch_device_kind"]
+private opaque deviceKindRaw (token : UInt32) : UInt32
+
 @[never_extract, extern "torchlean_libtorch_set_device"]
 private opaque setDeviceRaw (device : UInt32) : IO Unit
 
@@ -80,16 +89,30 @@ private opaque emptyCacheRaw (token : UInt32) : IO Unit
 @[no_expose] def version : IO String :=
   IO.lazyPure fun _ => versionRaw 0
 
-/-- Number of CUDA devices visible to the linked runtime; zero without LibTorch. -/
+/-- Number of CUDA devices visible to the linked runtime; zero without LibTorch and zero on a
+CPU-only SDK. A visible device is not necessarily the selected one: see `deviceKind`. -/
 @[no_expose] def deviceCount : IO UInt32 :=
   IO.lazyPure fun _ => deviceCountRaw 0
 
-/-- Device index selected by the bridge. This read alone does not establish availability. -/
+/-- The device the bridge runs on, as selected at its first call. -/
+inductive DeviceKind where
+  /-- A CUDA device. `Buffer.runtimeStatus` says whether one is usable. -/
+  | cuda
+  /-- The host: a CPU-only SDK, or `TORCHLEAN_LIBTORCH_DEVICE=cpu`. -/
+  | host
+  deriving DecidableEq, Repr
+
+/-- Which device the bridge selected. Meaningful only when LibTorch is linked: a build without it
+reads `.cuda` and `Buffer.runtimeStatus` reads `.notLinked`. -/
+@[no_expose] def deviceKind : IO DeviceKind :=
+  IO.lazyPure fun _ => if deviceKindRaw 0 == 0 then .host else .cuda
+
+/-- CUDA device index selected by the bridge. This read alone does not establish availability. -/
 @[no_expose] def getDevice : IO UInt32 :=
   IO.lazyPure fun _ => getDeviceRaw 0
 
 /--
-Select the device for subsequent bridge work.
+Select the CUDA device for subsequent bridge work; the host has no index and rejects the call.
 
 The native setter rejects changes while any buffer wrappers remain live, including released or
 empty wrappers. Retire those owners before switching devices; the call does not migrate tensors.
@@ -196,7 +219,7 @@ eligible choices causes a native execution error. Configure them before recordin
 Read the selected device's allocator memory fraction.
 
 This is a native allocator limit, not the fraction of memory currently free or a cache-only budget.
-Returns zero when LibTorch is not linked.
+Returns zero when LibTorch is not linked, and zero on the host, which has no such limit.
 -/
 @[no_expose] def getMemoryFraction : IO Float := do
   let fraction ← getMemoryFractionRaw 0
@@ -209,13 +232,15 @@ Set the selected device's allocator memory fraction to a finite value in `(0, 1]
 
 This configures native allocation policy; it does not release live tensors or reserve memory against
 other processes. Read back with `getMemoryFraction`; byte granularity may affect the observed limit.
+The host has no such limit and rejects the request.
 -/
 @[no_expose] def setMemoryFraction (fraction : Float) : IO Unit := do
   unless fraction.isFinite && 0.0 < fraction && fraction ≤ 1.0 do
     throw <| IO.userError "LibTorch: memory fraction must be finite and lie in (0, 1]"
   setMemoryFractionRaw fraction
 
-/-- Wait for the selected CUDA device's work, propagating native errors through `IO`. -/
+/-- Wait for the selected CUDA device's work, propagating native errors through `IO`. Host work
+is synchronous, so this returns at once there. -/
 @[no_expose] def synchronize : IO Unit :=
   synchronizeRaw 0
 
@@ -225,6 +250,7 @@ Release unused native allocator cache blocks.
 Live tensors, saved forward state, and library workspaces remain allocated. A cuBLAS workspace can
 outlive every TorchLean buffer, so both allocated and reserved bytes may remain after this call.
 Use the buffer ownership counters to distinguish those allocations from retained TorchLean owners.
+The host has no native cache, so this returns at once there.
 -/
 @[no_expose] def emptyCache : IO Unit :=
   emptyCacheRaw 0

NN/Runtime/Autograd/Engine/LibTorch/Trusted.lean

@@ -89,8 +89,9 @@ def Dtype.bytes : Dtype → Nat
   | .encoded => 0
 
 /--
-Opaque handle to a contiguous, dtype-carrying CUDA buffer in `csrc/libtorch/torchlean.cpp`.
-Builds without `-K cuda=true` cannot create one.
+Opaque handle to a contiguous, dtype-carrying buffer on the bridge's device (a CUDA device, or
+the host), implemented in `csrc/libtorch/torchlean.cpp`. Builds without `-K cuda=true` cannot
+create one.
 -/
 opaque BufferImpl : NonemptyType.{0}
 

NN/Tests/Runtime/Cuda/Elementwise.lean

@@ -735,13 +735,48 @@ private def checkCustomAutograd : IO Unit := do
   checkConfiguredAutograd
   checkRecordingBuffers
 
+/-- On the host a generated operation has no implementation: it must fail through `IO`, naming
+the requirement, before any CUDA call. Configured binary arithmetic, which only generated kernels
+compute, is refused when its tensor is uploaded, so a recording fails the same way. -/
+private def checkCustomComputationsOnHost : IO Unit := do
+  let input : Tensor Float32 [4] := [1, -2, 3, -4]
+  let outcome ← try
+    let _ ← square.run input (device := gpu)
+    pure (none : Option String)
+  catch error => pure (some (toString error))
+  match outcome with
+  | none => throw <| IO.userError "custom computation: a generated kernel ran on the host"
+  | some message =>
+      unless (message.splitOn "needs a CUDA device").length > 1 do
+        throw <| IO.userError s!"custom computation on the host: unexpected error: {message}"
+  let a : WideBinary := Rat.cast (1 + 1 / (2 ^ 100 : Nat) : Rat)
+  let configured : Tensor WideBinary [1] := Tensor.full [1] a
+  let recorded ← try
+    let _ ← (fun (x : WideBinary) => x * x).run configured (device := gpu) (grad := true)
+    pure (none : Option String)
+  catch error => pure (some (toString error))
+  match recorded with
+  | none => throw <| IO.userError "custom computation: configured arithmetic ran on the host"
+  | some message =>
+      unless (message.splitOn "needs a CUDA device").length > 1 do
+        throw <| IO.userError s!"configured arithmetic on the host: unexpected error: {message}"
+  IO.println "  custom computations: rejected on the host (generated kernels need a CUDA device)"
+
 /-- Exercise native elementwise VJPs, generated computations, and composed buffer operations. -/
 @[no_expose] def run : IO Unit := do
   IO.println "=== CUDA kernel coverage: elementwise ==="
   runActivationNumerics
   checkScaledProductExponential
-  checkCustomComputations
-  checkCustomAutograd
+  if (← Runtime.Autograd.LibTorch.deviceKind) == .cuda then
+    checkCustomComputations
+    checkCustomAutograd
+  else
+    checkCustomComputationsOnHost
+    -- Native-precision recordings and resident buffers need no generated kernel.
+    checkRecordedRecurrence 0
+    checkRecordedRecurrence 3
+    checkNativeAutograd
+    checkRecordingBuffers
 
   let s : Shape := [5]
   let a : Tensor Float s :=

NN/Tests/Runtime/Cuda/Stress.lean

@@ -779,6 +779,9 @@ def runMemoryTests : IO Unit := do
   | .nativeUnavailable =>
       throw <| IO.userError "LibTorch memory tests require a usable CUDA device"
   | .nativeAvailable => pure ()
+  if (← Runtime.Autograd.LibTorch.deviceKind) == .host then
+    IO.println "  skipped: the host device has no native allocator accounting"
+    return
   let self : System.FilePath := "/proc/self/exe"
   if !(← self.pathExists) then
     IO.println "  skipped: isolated memory tests require Linux /proc/self/exe"

NN/Tests/Suite.lean

@@ -52,7 +52,12 @@ def usage : String :=
     ]
 
 def run : IO Unit := do
+  -- `TORCHLEAN_REQUIRE_CUDA=1` means a GPU ran this: the bridge on the host does not satisfy it.
+  -- `TORCHLEAN_REQUIRE_NATIVE=1` means the bridge ran this, on whichever device it selected: the
+  -- check for a CPU-only SDK session, which must not pass on the unlinked default backend.
   if (← IO.getEnv "TORCHLEAN_REQUIRE_CUDA") == some "1" then
+    Runtime.Autograd.LibTorch.Buffer.requireCudaDevice
+  if (← IO.getEnv "TORCHLEAN_REQUIRE_NATIVE") == some "1" then
     Runtime.Autograd.LibTorch.Buffer.requireNativeRuntime
   -- Fresh subprocesses isolate native memory accounting and restore allocator limits on exit.
   -- These probes are selected before the ordinary suite to avoid unrelated live GPU owners.
@@ -80,11 +85,20 @@ def run : IO Unit := do
     | .notLinked =>
         IO.println "  CUDA kernels: skipped (LibTorch not linked)"
     | .nativeAvailable =>
+        -- The selected device must agree with what the bridge can see and what was asked for.
+        let kind ← Runtime.Autograd.LibTorch.deviceKind
+        let count ← Runtime.Autograd.LibTorch.deviceCount
+        let requested ← IO.getEnv "TORCHLEAN_LIBTORCH_DEVICE"
+        if kind == .cuda && count == 0 then
+          throw <| IO.userError "LibTorch selected a CUDA device while none is visible"
+        if requested == some "cpu" && kind != .host then
+          throw <| IO.userError "TORCHLEAN_LIBTORCH_DEVICE=cpu did not select the host"
+        IO.println s!"  LibTorch device: {match kind with | .cuda => "cuda" | .host => "host"}"
         NN.Tests.Runtime.EinsumDynamic.run .cuda
         Tests.Cuda.run
     | .nativeUnavailable =>
         throw <| IO.userError
-          "TorchLean was built with CUDA, but no usable CUDA device is visible"
+          "TorchLean was built with LibTorch, but no usable device is selected"
     IO.println "== TorchLean: all curated tests passed =="
 
 def main (args : List String) : IO Unit := do

README.md

@@ -34,7 +34,8 @@ TorchLean's backend architecture, see the [Installation guide](https://lean-dojo
 scripts/lake.sh exe torchlean quickstart_mlp --device cpu --steps 10 --arithmetic ieee --execution eager
 scripts/lake.sh exe torchlean quickstart_mlp --device cpu --steps 10 --execution eager
 
-# Optional GPU run with a CUDA-enabled LibTorch SDK, matching toolkit, and NVIDIA GPU:
+# Optional GPU run with a CUDA-enabled LibTorch SDK, matching toolkit, and NVIDIA GPU
+# (a CPU-only LibTorch SDK runs the same backend on the host; see csrc/libtorch/README.md):
 export TORCHLEAN_LIBTORCH_HOME=/absolute/path/to/libtorch
 scripts/lake.sh -Kcuda=true build
 scripts/lake.sh -Kcuda=true exe torchlean quickstart_mlp --device cuda --steps 10 --execution eager