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

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

@@ -10,10 +10,11 @@ public import NN.Runtime.Autograd.Engine.LibTorch.Controls
 public import NN.Runtime.Autograd.Engine.LibTorch.Trusted
 
 /-!
-# CUDA Float32 Buffers
+# LibTorch Float32 Buffers
 
-Low-level float32 tensor 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 float32 tensor 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.
 -/
@@ -28,28 +29,30 @@ namespace 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, for a user-selected `cuda` session. -/
 def requireNativeRuntime : IO Unit :=
   match runtimeStatus with
   | .nativeAvailable => pure ()
@@ -59,7 +62,8 @@ def requireNativeRuntime : IO Unit :=
           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"
+        "CUDA was requested and this is a LibTorch build, but no usable CUDA device is visible; \
+          TORCHLEAN_LIBTORCH_DEVICE=cpu selects the host instead"
 
 /-! ### 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

@@ -74,8 +74,8 @@ namespace Autograd
 namespace LibTorch
 
 /--
-Opaque handle to a contiguous float32 CUDA buffer, implemented in `csrc/libtorch/torchlean.cpp`.
-Builds without `-K cuda=true` cannot create one.
+Opaque handle to a contiguous float32 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/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

@@ -167,12 +167,21 @@ 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.API.BufferUpdates.checkStochasticBuffers (device := .cuda)
         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

@@ -28,7 +28,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

csrc/libtorch/CMakeLists.txt

@@ -2,7 +2,7 @@ cmake_minimum_required(VERSION 3.22)
 project(TorchLeanLibTorch LANGUAGES CXX)
 
 if(NOT CMAKE_SYSTEM_NAME STREQUAL "Linux")
-  message(FATAL_ERROR "TorchLean's LibTorch CUDA backend currently requires Linux.")
+  message(FATAL_ERROR "TorchLean's LibTorch backend currently requires Linux.")
 endif()
 if(NOT EXISTS "${TORCHLEAN_LEAN_INCLUDE}/lean/lean.h")
   message(FATAL_ERROR "TORCHLEAN_LEAN_INCLUDE must name the pinned Lean include directory.")
@@ -16,8 +16,12 @@ set(Torch_DIR "${TORCHLEAN_LIBTORCH_HOME}/share/cmake/Torch")
 set(Caffe2_DIR "${TORCHLEAN_LIBTORCH_HOME}/share/cmake/Caffe2")
 list(PREPEND CMAKE_PREFIX_PATH "${TORCHLEAN_LIBTORCH_HOME}")
 find_package(Torch REQUIRED CONFIG PATHS "${Torch_DIR}" NO_DEFAULT_PATH)
-if(NOT TARGET torch_cuda)
-  message(FATAL_ERROR "cuda=true requires a CUDA-enabled LibTorch SDK (torch_cuda is missing).")
+# A CUDA-enabled SDK exports `torch_cuda`; a CPU-only SDK does not, and the backend then runs on
+# the host. The source reads the choice as TORCHLEAN_LIBTORCH_CUDA.
+if(TARGET torch_cuda)
+  set(TORCHLEAN_LIBTORCH_CUDA 1)
+else()
+  set(TORCHLEAN_LIBTORCH_CUDA 0)
 endif()
 
 get_target_property(torch_cxx_standard torch CXX_STANDARD)
@@ -54,17 +58,25 @@ endforeach()
 # This executable is built as a dependency of the backend and is never run.
 file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/sdk_link_check.cpp" [=[
   #include <ATen/ATen.h>
+  #if TORCHLEAN_LIBTORCH_CUDA
   #include <c10/cuda/CUDAFunctions.h>
+  #endif
   #include <torch/version.h>
   #include <string>
   int main() {
     const c10::Device device(std::string("cpu"));
     auto tensor = at::zeros({1}, at::TensorOptions().device(device));
+  #if TORCHLEAN_LIBTORCH_CUDA
     return tensor.numel() == 1 && c10::cuda::device_count() >= 0 ? 0 : 1;
+  #else
+    return tensor.numel() == 1 ? 0 : 1;
+  #endif
   }
 ]=])
 add_executable(torchlean_sdk_link_check EXCLUDE_FROM_ALL
   "${CMAKE_CURRENT_BINARY_DIR}/sdk_link_check.cpp")
+target_compile_definitions(torchlean_sdk_link_check PRIVATE
+  TORCHLEAN_LIBTORCH_CUDA=${TORCHLEAN_LIBTORCH_CUDA})
 target_link_libraries(torchlean_sdk_link_check PRIVATE ${TORCH_LIBRARIES})
 separate_arguments(torch_sdk_cxx_flags NATIVE_COMMAND "${TORCH_CXX_FLAGS}")
 target_compile_options(torchlean_sdk_link_check PRIVATE ${torch_sdk_cxx_flags})
@@ -74,7 +86,8 @@ add_library(torchlean_libtorch SHARED
   torchlean.cpp
 )
 add_dependencies(torchlean_libtorch torchlean_sdk_link_check)
-target_compile_definitions(torchlean_libtorch PRIVATE TORCHLEAN_LIBTORCH)
+target_compile_definitions(torchlean_libtorch PRIVATE
+  TORCHLEAN_LIBTORCH TORCHLEAN_LIBTORCH_CUDA=${TORCHLEAN_LIBTORCH_CUDA})
 target_compile_options(torchlean_libtorch PRIVATE ${torch_sdk_cxx_flags})
 target_include_directories(torchlean_libtorch PRIVATE
   "${TORCHLEAN_LEAN_INCLUDE}"
@@ -96,6 +109,7 @@ target_link_options(torchlean_libtorch PRIVATE "LINKER:--disable-new-dtags")
 file(WRITE "${CMAKE_BINARY_DIR}/sdk.txt"
   "Torch_VERSION=${torch_sdk_version}
 Torch_DIR=${Torch_DIR}
+TORCHLEAN_LIBTORCH_CUDA=${TORCHLEAN_LIBTORCH_CUDA}
 TORCH_CXX_FLAGS=${TORCH_CXX_FLAGS}
 TORCH_COMPILE_FEATURES=${torch_compile_features}
 TORCH_COMPILE_OPTIONS=${torch_compile_options}
@@ -106,5 +120,5 @@ CXX_COMPILER=${CMAKE_CXX_COMPILER}
 CXX_COMPILER_ID=${CMAKE_CXX_COMPILER_ID}
 CXX_COMPILER_VERSION=${CMAKE_CXX_COMPILER_VERSION}
 ")
-message(STATUS "TorchLean: LibTorch ${torch_sdk_version}, C++${CMAKE_CXX_STANDARD}")
+message(STATUS "TorchLean: LibTorch ${torch_sdk_version}, C++${CMAKE_CXX_STANDARD}, CUDA support ${TORCHLEAN_LIBTORCH_CUDA}")
 message(STATUS "TorchLean SDK flags: ${TORCH_CXX_FLAGS}; ${torch_compile_options}")

csrc/libtorch/README.md

@@ -1,8 +1,10 @@
 # TorchLean LibTorch backend
 
-This directory holds the CUDA backend behind TorchLean's GPU buffer ABI. It calls the selected
-LibTorch SDK's ATen operations. Lean checks shapes and dispatches through the extern symbols;
-native memory safety, SDK behavior, and floating-point execution remain outside Lean's kernel.
+This directory holds the LibTorch backend behind TorchLean's device buffer ABI. It calls the
+selected LibTorch SDK's ATen operations on a CUDA device, or on the host when the SDK is CPU-only
+or `TORCHLEAN_LIBTORCH_DEVICE=cpu` says so. Lean checks shapes and dispatches through the extern
+symbols; native memory safety, SDK behavior, and floating-point execution remain outside Lean's
+kernel.
 
 ## Layout
 
@@ -23,7 +25,7 @@ compiled executables.
 ## Build selection
 
 `scripts/lake.sh build` selects the default `pureLean`/`portableCPU` build, without an SDK or
-toolkit. `cuda=true` requires the complete LibTorch CUDA backend.
+toolkit. `cuda=true` links the LibTorch backend; the SDK it is built against decides the device.
 
 A full SDK contains `include/`,
 `lib/`, and `share/cmake/Torch/TorchConfig.cmake`; a partial header snapshot is insufficient.
@@ -40,9 +42,36 @@ scripts/checks/check.sh --libtorch-home "$TORCHLEAN_LIBTORCH_HOME" --ci-all
 under the package root. The build requires Linux, CMake 3.22 or newer, Make, the pinned Lean
 headers, and a compatible C++20 compiler. SDK CMake discovers the ABI, any stricter C++ standard,
 transitive libraries, and rpath. An executable built in the same project checks compiler/link
-compatibility without running. SDK discovery may require a matching CUDA
+compatibility without running. A CUDA-enabled SDK's discovery may require a matching CUDA
 development toolkit, even though TorchLean itself compiles only C++ sources.
 
+### Host device
+
+A CPU-only SDK builds the same backend without CUDA support (CMake sees no `torch_cuda` target
+and compiles with `TORCHLEAN_LIBTORCH_CUDA=0`), and the bridge then runs every operation on the
+host through ATen's CPU kernels. No CUDA toolkit or GPU is involved, so the curated suite
+exercises the real buffer ABI on a hosted runner:
+
+```bash
+python3 -m venv venv-cpu
+venv-cpu/bin/pip install --index-url https://download.pytorch.org/whl/cpu torch==2.11.0
+export TORCHLEAN_LIBTORCH_HOME="$PWD/venv-cpu/lib/python3.12/site-packages/torch"
+scripts/lake.sh -Kcuda=true build NN NNCI NNExamples NNTests nn_tests_suite
+scripts/lake.sh -Kcuda=true test
+```
+
+The device is selected once, at the bridge's first call, and `Runtime.Autograd.LibTorch.deviceKind`
+reports it. With a CUDA-enabled SDK, `TORCHLEAN_LIBTORCH_DEVICE=cpu` selects the host on a machine
+that also has a GPU (the suite then runs on the host), and `TORCHLEAN_LIBTORCH_DEVICE=cuda` insists
+on a CUDA device. A CUDA-enabled SDK without a visible device reports `RuntimeStatus.nativeUnavailable`
+rather than falling back to the host. On the host, the native allocator counters and driver memory
+read zero, `synchronize` and `emptyCache` return at once, `setMemoryFraction` and `setDevice` are
+rejected, and the memory accounting probes of the suite are skipped. Results are bit-identical
+across the two devices only for programs made of the IEEE basic operations, gathers and lookups:
+ATen's CPU and CUDA transcendental functions (`exp`, `log`, `tanh`, ...) come from different
+libraries and differ at the ulp level, and reductions differ in order. The suite's tolerances
+hold on both.
+
 ## Tested SDK versions
 
 The current adapter compiled and passed the curated CUDA suite and focused attention check on