From 9fcddfb8bc58248445af3dae2dea1ce60b4268ca Mon Sep 17 00:00:00 2001 From: Nicolas Rouquette Date: Sat, 3 Oct 2026 17:29:28 -0700 Subject: [PATCH 1/2] feat(libtorch): run the bridge on the host when the SDK is CPU-only The LibTorch bridge pinned `c10::kCUDA`, so a build without a CUDA device had no backend at all: `unavailable.c` fails every buffer operation, and the device buffer suite could only run on a GPU host. The bridge now selects its device once, at its first call, and runs on the host through ATen's CPU kernels when the SDK has no CUDA support or when `TORCHLEAN_LIBTORCH_DEVICE=cpu` asks for it. Nothing is reimplemented: every operation already goes through `torchlean::options()`. Selection: `TORCHLEAN_LIBTORCH_DEVICE=cpu` selects the host and `=cuda` insists on a CUDA device. Unset, a visible CUDA device is selected when the SDK has CUDA support; otherwise the host is selected only when the SDK has no CUDA support at all, so a CUDA-enabled SDK without a visible device stays `RuntimeStatus.nativeUnavailable` rather than silently running on the host. CMake no longer requires the `torch_cuda` target. A CPU-only SDK compiles the backend with `TORCHLEAN_LIBTORCH_CUDA=0`, without CUDA headers or a toolkit; the CPU-only pip wheel is a complete SDK for it. `sdk.txt` records the choice. Lean: `Runtime.Autograd.LibTorch.DeviceKind` and `deviceKind` report the selection. `RuntimeStatus` keeps its three states, with `.nativeAvailable` meaning that a device is selected, CUDA or host, so `requireNativeRuntime` and every existing match are unchanged. 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 suite skips its native memory probes. The suite prints the device and checks that the selection agrees with the visible device count and with the environment variable. Verified with pip torch 2.11.0+cu128 on an RTX A4500 (device cuda) and with the same binary under `TORCHLEAN_LIBTORCH_DEVICE=cpu` (device host), and with pip torch 2.11.0+cpu and no toolkit (`TORCHLEAN_LIBTORCH_CUDA=0`, device host): all curated tests passed in each run, with only the three memory probes skipped on the host. The default build and its CPU suite are unchanged. --- .../Autograd/Engine/LibTorch/Buffer.lean | 22 +-- .../Autograd/Engine/LibTorch/Controls.lean | 38 +++++- .../Autograd/Engine/LibTorch/Trusted.lean | 4 +- NN/Tests/Runtime/Cuda/Stress.lean | 3 + NN/Tests/Suite.lean | 11 +- README.md | 3 +- csrc/libtorch/CMakeLists.txt | 24 +++- csrc/libtorch/README.md | 37 ++++- csrc/libtorch/torchlean.cpp | 126 +++++++++++++++--- csrc/libtorch/torchlean_libtorch.h | 7 +- csrc/libtorch/unavailable.c | 9 +- lakefile.lean | 13 +- scripts/libtorch_build.py | 6 +- 13 files changed, 246 insertions(+), 57 deletions(-) diff --git a/NN/Runtime/Autograd/Engine/LibTorch/Buffer.lean b/NN/Runtime/Autograd/Engine/LibTorch/Buffer.lean index b8bbd092..68bb6e8a 100644 --- a/NN/Runtime/Autograd/Engine/LibTorch/Buffer.lean +++ b/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,13 +29,15 @@ 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 @@ -42,14 +45,14 @@ inductive RuntimeStatus where @[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 -/ diff --git a/NN/Runtime/Autograd/Engine/LibTorch/Controls.lean b/NN/Runtime/Autograd/Engine/LibTorch/Controls.lean index 2fbcec36..e3fd9246 100644 --- a/NN/Runtime/Autograd/Engine/LibTorch/Controls.lean +++ b/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 diff --git a/NN/Runtime/Autograd/Engine/LibTorch/Trusted.lean b/NN/Runtime/Autograd/Engine/LibTorch/Trusted.lean index 69abe3aa..d5acc650 100644 --- a/NN/Runtime/Autograd/Engine/LibTorch/Trusted.lean +++ b/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} diff --git a/NN/Tests/Runtime/Cuda/Stress.lean b/NN/Tests/Runtime/Cuda/Stress.lean index eed49071..cecf57e5 100644 --- a/NN/Tests/Runtime/Cuda/Stress.lean +++ b/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" diff --git a/NN/Tests/Suite.lean b/NN/Tests/Suite.lean index fb4f68ec..d310c371 100644 --- a/NN/Tests/Suite.lean +++ b/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 diff --git a/README.md b/README.md index a17be1b9..da8f02fa 100644 --- a/README.md +++ b/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 diff --git a/csrc/libtorch/CMakeLists.txt b/csrc/libtorch/CMakeLists.txt index b950ec48..f538e109 100644 --- a/csrc/libtorch/CMakeLists.txt +++ b/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 + #if TORCHLEAN_LIBTORCH_CUDA #include + #endif #include #include 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}") diff --git a/csrc/libtorch/README.md b/csrc/libtorch/README.md index 56130682..5be73b7a 100644 --- a/csrc/libtorch/README.md +++ b/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,34 @@ 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. ATen's CPU and CUDA kernels +agree on IEEE elementwise float32 but not on reduction order, so results of reductions are not +bit-identical across the two devices. + ## Tested SDK versions The current adapter compiled and passed the curated CUDA suite and focused attention check on diff --git a/csrc/libtorch/torchlean.cpp b/csrc/libtorch/torchlean.cpp index 09ffe5bc..b5bc84bb 100644 --- a/csrc/libtorch/torchlean.cpp +++ b/csrc/libtorch/torchlean.cpp @@ -3,18 +3,28 @@ // Buffer ownership and runtime controls +// TORCHLEAN_LIBTORCH_CUDA is 1 for an SDK that exports `torch_cuda` and 0 for a CPU-only SDK; +// CMake sets it from the selected SDK. The bridge then runs on a CUDA device or on the host. +#ifndef TORCHLEAN_LIBTORCH_CUDA +#define TORCHLEAN_LIBTORCH_CUDA 1 +#endif + #include +#if TORCHLEAN_LIBTORCH_CUDA #include #include +#endif #include #include #include #include +#include #include #include #include #include +#include namespace { @@ -42,6 +52,39 @@ Counter wrappers; std::once_flag initialization; std::atomic selected_device{0}; +// Where the bridge runs, decided once by `torchlean::initialize` (see `select_device`): a CUDA +// device, the host, or nothing usable. +enum class Selection { cuda, host, unavailable }; +Selection selection = Selection::unavailable; + +uint32_t cuda_device_count() { +#if TORCHLEAN_LIBTORCH_CUDA + return static_cast(c10::cuda::device_count()); +#else + return 0; +#endif +} + +// `TORCHLEAN_LIBTORCH_DEVICE` selects "cuda" or "cpu" explicitly. Unset, a visible CUDA device +// is selected when the SDK has CUDA support; otherwise the host is selected only when the SDK +// has no CUDA support at all. A CUDA SDK without a visible device stays unavailable rather than +// silently running on the host. +Selection select_device() { + const char* requested = std::getenv("TORCHLEAN_LIBTORCH_DEVICE"); + const std::string value = requested ? requested : ""; + if (value == "cpu") return Selection::host; + if (value == "cuda") return cuda_device_count() > 0 ? Selection::cuda : Selection::unavailable; + if (!value.empty()) { + std::fprintf(stderr, + "LibTorch: TORCHLEAN_LIBTORCH_DEVICE=\"%s\" is neither \"cuda\" nor \"cpu\"; " + "no device selected\n", + value.c_str()); + return Selection::unavailable; + } + if (cuda_device_count() > 0) return Selection::cuda; + return TORCHLEAN_LIBTORCH_CUDA ? Selection::unavailable : Selection::host; +} + bool release_data(torchlean_cuda_buffer* buffer) { if (!buffer || !buffer->tensor.defined()) return false; const size_t size = buffer->size; @@ -196,16 +239,29 @@ lean_obj_res download_bytes(b_lean_obj_arg object) { return out; } -auto allocator_stats() { - return c10::cuda::CUDACachingAllocator::getDeviceStats(torchlean::device().index()); +// Native allocator and driver counters exist for the CUDA device only; the host reads zero. +#if TORCHLEAN_LIBTORCH_CUDA +template +uint64_t allocator_stat(F&& field) { + return torchlean::invoke([&]() -> uint64_t { + if (selection != Selection::cuda) return 0; + return field(c10::cuda::CUDACachingAllocator::getDeviceStats(torchlean::device().index())); + }); } +#endif uint64_t memory_info(bool total) { +#if TORCHLEAN_LIBTORCH_CUDA return torchlean::invoke([&]() -> uint64_t { + if (selection != Selection::cuda) return 0; size_t free_bytes = 0, total_bytes = 0; C10_CUDA_CHECK(cudaMemGetInfo(&free_bytes, &total_bytes)); return total ? total_bytes : free_bytes; }); +#else + (void)total; + return 0; +#endif } // Initial policy from TORCHLEAN_CUDA_DETERMINISTIC_REDUCTIONS: unset, empty, or "0" means off. @@ -232,11 +288,13 @@ void initialize() { const bool deterministic = deterministic_reductions_requested(); context.setDeterministicAlgorithms(deterministic, false); context.setDeterministicCuDNN(deterministic); - context.lazyInitDevice(c10::kCUDA); + selection = select_device(); + if (selection == Selection::cuda) context.lazyInitDevice(c10::kCUDA); }); } c10::Device device() { + if (selection == Selection::host) return c10::Device(c10::kCPU); return c10::Device(c10::kCUDA, static_cast(selected_device.load())); } @@ -247,8 +305,9 @@ at::TensorOptions options() { const at::Tensor& tensor(b_lean_obj_arg object) { const auto* buffer = torchlean_cuda_buffer_unbox(object); TORCH_CHECK(buffer->tensor.defined(), "LibTorch: buffer has been released"); - TORCH_CHECK(buffer->tensor.is_cuda() && buffer->tensor.scalar_type() == at::kFloat, - "LibTorch: expected a CUDA float32 buffer"); + TORCH_CHECK(buffer->tensor.device().type() == device().type() && + buffer->tensor.scalar_type() == at::kFloat, + "LibTorch: expected a float32 buffer on the selected device"); TORCH_CHECK(!buffer->tensor.requires_grad(), "LibTorch: unexpected autograd tensor"); TORCH_CHECK(buffer->tensor.numel() == static_cast(buffer->size), "LibTorch: buffer size disagrees with storage"); @@ -256,8 +315,9 @@ const at::Tensor& tensor(b_lean_obj_arg object) { } torchlean_cuda_buffer* owned(at::Tensor value) { - TORCH_CHECK(value.defined() && value.is_cuda() && value.scalar_type() == at::kFloat, - "LibTorch: expected a CUDA float32 result"); + TORCH_CHECK(value.defined() && value.device().type() == device().type() && + value.scalar_type() == at::kFloat, + "LibTorch: expected a float32 result on the selected device"); TORCH_CHECK(!value.requires_grad(), "LibTorch: native operations must not record autograd"); auto buffer = std::make_unique(); buffer->tensor = value.reshape({-1}).contiguous(); @@ -299,8 +359,24 @@ extern "C" void torchlean_cuda_buffer_drop_unboxed(torchlean_cuda_buffer* buffer delete buffer; } +// 1 = a device is selected (CUDA, or the host), 2 = nothing usable; `unavailable.c` reports 0. extern "C" LEAN_EXPORT uint32_t torchlean_cuda_runtime_status(uint32_t) { - return c10::cuda::device_count() > 0 ? 1 : 2; + try { + torchlean::initialize(); + } catch (const std::exception&) { + return 2; + } + return selection == Selection::unavailable ? 2 : 1; +} + +// 0 = the host, 1 = CUDA (a CUDA device is selected, or none is usable on a CUDA SDK). +extern "C" LEAN_EXPORT uint32_t torchlean_libtorch_device_kind(uint32_t) { + try { + torchlean::initialize(); + } catch (const std::exception&) { + return 1; + } + return selection == Selection::host ? 0 : 1; } #define TORCHLEAN_COUNTER(NAME, VALUE) \ @@ -325,10 +401,15 @@ extern "C" LEAN_EXPORT uint64_t torchlean_cuda_allocator_device_total_bytes(uint return memory_info(true); } +#if TORCHLEAN_LIBTORCH_CUDA #define TORCHLEAN_ALLOCATOR_STAT(NAME, FIELD, WHICH) \ extern "C" LEAN_EXPORT uint64_t torchlean_libtorch_##NAME(uint32_t) { \ - return torchlean::invoke([]() -> uint64_t { return allocator_stats().FIELD[0].WHICH; }); \ + return allocator_stat([](const auto& stats) -> uint64_t { return stats.FIELD[0].WHICH; }); \ } +#else +#define TORCHLEAN_ALLOCATOR_STAT(NAME, FIELD, WHICH) \ + extern "C" LEAN_EXPORT uint64_t torchlean_libtorch_##NAME(uint32_t) { return 0; } +#endif TORCHLEAN_ALLOCATOR_STAT(allocated_bytes, allocated_bytes, current) TORCHLEAN_ALLOCATOR_STAT(reserved_bytes, reserved_bytes, current) TORCHLEAN_ALLOCATOR_STAT(peak_allocated_bytes, allocated_bytes, peak) @@ -408,7 +489,7 @@ extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_version(uint32_t) { } extern "C" LEAN_EXPORT uint32_t torchlean_libtorch_device_count(uint32_t) { - return static_cast(c10::cuda::device_count()); + return cuda_device_count(); } extern "C" LEAN_EXPORT uint32_t torchlean_libtorch_get_device(uint32_t) { @@ -417,8 +498,8 @@ extern "C" LEAN_EXPORT uint32_t torchlean_libtorch_get_device(uint32_t) { extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_set_device(uint32_t index) { return io([&] { - TORCH_CHECK(index < static_cast(c10::cuda::device_count()), - "LibTorch: selected CUDA device does not exist"); + TORCH_CHECK(selection == Selection::cuda, "LibTorch: no CUDA device is selected"); + TORCH_CHECK(index < cuda_device_count(), "LibTorch: selected CUDA device does not exist"); TORCH_CHECK(wrappers.live.load() == 0, "LibTorch: select the device before creating tensor buffers"); selected_device.store(static_cast(index)); @@ -480,8 +561,12 @@ extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_set_setting( extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_get_memory_fraction(uint32_t) { return io([] { - return lean_box_float( - c10::cuda::CUDACachingAllocator::getMemoryFraction(torchlean::device().index())); + double fraction = 0.0; +#if TORCHLEAN_LIBTORCH_CUDA + if (selection == Selection::cuda) + fraction = c10::cuda::CUDACachingAllocator::getMemoryFraction(torchlean::device().index()); +#endif + return lean_box_float(fraction); }); } @@ -489,21 +574,30 @@ extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_set_memory_fraction(doubl return io([&] { TORCH_CHECK(std::isfinite(fraction) && fraction > 0.0 && fraction <= 1.0, "LibTorch: memory fraction must be finite and in (0, 1]"); + TORCH_CHECK(selection == Selection::cuda, + "LibTorch: the allocator memory fraction applies to a CUDA device only"); +#if TORCHLEAN_LIBTORCH_CUDA c10::cuda::CUDACachingAllocator::setMemoryFraction(fraction, torchlean::device().index()); +#endif return lean_box(0); }); } +// Host work is synchronous and has no native cache to release: both are no-ops there. extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_synchronize(uint32_t) { return io([] { - c10::cuda::device_synchronize(); +#if TORCHLEAN_LIBTORCH_CUDA + if (selection == Selection::cuda) c10::cuda::device_synchronize(); +#endif return lean_box(0); }); } extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_empty_cache(uint32_t) { return io([] { - c10::cuda::CUDACachingAllocator::emptyCache(); +#if TORCHLEAN_LIBTORCH_CUDA + if (selection == Selection::cuda) c10::cuda::CUDACachingAllocator::emptyCache(); +#endif return lean_box(0); }); } diff --git a/csrc/libtorch/torchlean_libtorch.h b/csrc/libtorch/torchlean_libtorch.h index 2e216228..bff22540 100644 --- a/csrc/libtorch/torchlean_libtorch.h +++ b/csrc/libtorch/torchlean_libtorch.h @@ -17,9 +17,10 @@ #include // Native side of `NN.Runtime.Autograd.Engine.LibTorch.Buffer`. Lean owns an external object that points -// at a `torchlean_cuda_buffer`; `size` counts float32 elements, not bytes. Callers validate shape -// metadata before touching storage. This is a trusted boundary: Lean proves shape contracts around -// these calls but cannot see tensor lifetimes or CUDA behavior. +// at a `torchlean_cuda_buffer`; `size` counts float32 elements, not bytes, and the tensor lives on +// the device the bridge selected (a CUDA device, or the host). Callers validate shape metadata +// before touching storage. This is a trusted boundary: Lean proves shape contracts around these +// calls but cannot see tensor lifetimes or device behavior. struct torchlean_cuda_buffer { size_t size; at::Tensor tensor; diff --git a/csrc/libtorch/unavailable.c b/csrc/libtorch/unavailable.c index 49c9a668..e3a6dcaa 100644 --- a/csrc/libtorch/unavailable.c +++ b/csrc/libtorch/unavailable.c @@ -18,12 +18,19 @@ static lean_obj_res unavailable_io(void) { lean_mk_io_user_error(lean_mk_string(TORCHLEAN_UNAVAILABLE_MESSAGE))); } -// 0 = built without LibTorch, 1 = LibTorch with a visible device, 2 = LibTorch without one. +// 0 = built without LibTorch, 1 = LibTorch with a selected device (CUDA, or the host), +// 2 = LibTorch without one. LEAN_EXPORT uint32_t torchlean_cuda_runtime_status(uint32_t token) { (void)token; return 0u; } +// The device kind (0 = host, 1 = CUDA) means nothing without LibTorch; `runtimeStatus` says so. +LEAN_EXPORT uint32_t torchlean_libtorch_device_kind(uint32_t token) { + (void)token; + return 1u; +} + LEAN_EXPORT lean_obj_res torchlean_libtorch_version(uint32_t token) { (void)token; return lean_mk_string("unavailable"); diff --git a/lakefile.lean b/lakefile.lean index 5dc92c79..4683a120 100644 --- a/lakefile.lean +++ b/lakefile.lean @@ -9,7 +9,8 @@ import Lake.Util.Proc open Lake DSL open System -/-- Whether Lake should link the LibTorch CUDA backend instead of the unavailable-backend shim. -/ +/-- Whether Lake should link the LibTorch backend instead of the unavailable-backend shim. The +SDK decides the device: CUDA-enabled SDKs run on a CUDA device, CPU-only SDKs on the host. -/ private def cudaEnabled : Bool := let value := (get_config? cuda).getD "false" value == "true" || value == "1" @@ -51,9 +52,10 @@ package TorchLean where ## Native backend libraries `-K cuda=true` builds one LibTorch C++ library containing the numerical C ABI exports. CMake obtains -the ABI, language standard, libraries, and runtime paths from the selected SDK. The default build -needs neither LibTorch nor a CUDA toolkit: it links a small C file that reports the backend as not -linked and fails every GPU call with an explanation. +the ABI, language standard, libraries, and runtime paths from the selected SDK; a CPU-only SDK +yields a backend that runs on the host, a CUDA-enabled one a backend that runs on a CUDA device. +The default build needs neither LibTorch nor a CUDA toolkit: it links a small C file that reports +the backend as not linked and fails every GPU call with an explanation. -/ /-- @@ -86,7 +88,8 @@ private def nativeCompilerJob (name : String) : SpawnM (Job FilePath) := Job.asy traceNativeTool compiler return compiler -/-- Numerical CUDA primitives built and linked with the selected LibTorch SDK. -/ +/-- Numerical primitives built and linked with the selected LibTorch SDK, CUDA-enabled or +CPU-only. -/ target torchlean_libtorch pkg : FilePath := do let lean ← getLeanInstall let scriptJob ← inputFile (pkg.dir / "scripts/libtorch_build.py") false diff --git a/scripts/libtorch_build.py b/scripts/libtorch_build.py index 664c249a..d5a113ea 100755 --- a/scripts/libtorch_build.py +++ b/scripts/libtorch_build.py @@ -297,9 +297,9 @@ def main() -> int: print(home if args.resolve_home else build(args, package, home)) except (ValueError, OSError, subprocess.CalledProcessError) as error: print(f"error: LibTorch build: {error}", file=sys.stderr) - print("cuda=true requires a compatible LibTorch CUDA SDK. " - "SDK CMake discovery may require a matching CUDA development toolkit.", - file=sys.stderr) + print("cuda=true requires a compatible LibTorch SDK: CUDA-enabled, or CPU-only for the " + "host device. A CUDA-enabled SDK's CMake discovery may require a matching CUDA " + "development toolkit.", file=sys.stderr) return 1 return 0 From bed0d9bbff32f023e74e3c0ff8191e69d361f52a Mon Sep 17 00:00:00 2001 From: Nicolas Rouquette Date: Sat, 3 Oct 2026 18:25:29 -0700 Subject: [PATCH 2/2] docs(libtorch): say which results are bit-identical across the host and CUDA Programs made of the IEEE basic operations, gathers and lookups give bit-identical results on the two devices; transcendental functions come from different libraries on the host and on CUDA and differ at the ulp level, and reductions differ in order. --- csrc/libtorch/README.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/csrc/libtorch/README.md b/csrc/libtorch/README.md index 5be73b7a..299b7969 100644 --- a/csrc/libtorch/README.md +++ b/csrc/libtorch/README.md @@ -66,9 +66,11 @@ that also has a GPU (the suite then runs on the host), and `TORCHLEAN_LIBTORCH_D 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. ATen's CPU and CUDA kernels -agree on IEEE elementwise float32 but not on reduction order, so results of reductions are not -bit-identical across the two devices. +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