From a3a86749448630f001c093662af0e78b3a861607 Mon Sep 17 00:00:00 2001 From: Nicolas Rouquette Date: Sat, 3 Oct 2026 17:29:28 -0700 Subject: [PATCH 1/5] 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 ATen 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, the CUDA toolkit, NVRTC or the driver API; the CPU-only pip wheel is a complete SDK for it. `sdk.txt` records the choice. Runtime-compiled kernels (the generated operations of `NN.Kernel`) need NVRTC and the driver API, so they compile only with a CUDA-enabled SDK, and `run_kernel` refuses at its entry unless a CUDA device is selected: on the host a generated operation fails through `IO` with an error that names the requirement, before any CUDA call. The CUDA suite runs its custom-computation checks on a CUDA device and checks that refusal on the host. 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. --- NN/Kernel/Runtime.lean | 5 +- .../Autograd/Engine/LibTorch/Buffer.lean | 22 +-- .../Autograd/Engine/LibTorch/Controls.lean | 38 ++++- .../Autograd/Engine/LibTorch/Trusted.lean | 5 +- NN/Tests/Runtime/Cuda/Elementwise.lean | 22 ++- NN/Tests/Runtime/Cuda/Stress.lean | 3 + NN/Tests/Suite.lean | 11 +- README.md | 3 +- csrc/libtorch/CMakeLists.txt | 33 +++- csrc/libtorch/README.md | 38 ++++- csrc/libtorch/torchlean.cpp | 144 +++++++++++++++--- csrc/libtorch/torchlean_libtorch.h | 7 +- csrc/libtorch/unavailable.c | 9 +- lakefile.lean | 14 +- scripts/README.md | 9 +- scripts/libtorch_build.py | 6 +- 16 files changed, 300 insertions(+), 69 deletions(-) diff --git a/NN/Kernel/Runtime.lean b/NN/Kernel/Runtime.lean index d000c78f4..7ca26874e 100644 --- a/NN/Kernel/Runtime.lean +++ b/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 diff --git a/NN/Runtime/Autograd/Engine/LibTorch/Buffer.lean b/NN/Runtime/Autograd/Engine/LibTorch/Buffer.lean index dda85b720..d816fe937 100644 --- a/NN/Runtime/Autograd/Engine/LibTorch/Buffer.lean +++ b/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,13 +55,15 @@ 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 @@ -68,14 +71,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 () @@ -85,7 +88,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 2fbcec36e..e3fd9246b 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 7a64b176b..b332a956b 100644 --- a/NN/Runtime/Autograd/Engine/LibTorch/Trusted.lean +++ b/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} diff --git a/NN/Tests/Runtime/Cuda/Elementwise.lean b/NN/Tests/Runtime/Cuda/Elementwise.lean index 4193bb275..475845d2b 100644 --- a/NN/Tests/Runtime/Cuda/Elementwise.lean +++ b/NN/Tests/Runtime/Cuda/Elementwise.lean @@ -735,13 +735,31 @@ 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. -/ +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}" + 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 let s : Shape := [5] let a : Tensor Float s := diff --git a/NN/Tests/Runtime/Cuda/Stress.lean b/NN/Tests/Runtime/Cuda/Stress.lean index 7cef8859f..7569060c9 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 77bed1efc..0cff42880 100644 --- a/NN/Tests/Suite.lean +++ b/NN/Tests/Suite.lean @@ -80,11 +80,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 diff --git a/README.md b/README.md index 3cbf96879..794dd752a 100644 --- a/README.md +++ b/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 diff --git a/csrc/libtorch/CMakeLists.txt b/csrc/libtorch/CMakeLists.txt index 6e9d2ade7..69ccf4694 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,9 +16,14 @@ 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) -find_package(CUDAToolkit REQUIRED) -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. Runtime-compiled kernels +# need NVRTC and the driver API, so only a CUDA-enabled SDK looks for the CUDA toolkit. +if(TARGET torch_cuda) + set(TORCHLEAN_LIBTORCH_CUDA 1) + find_package(CUDAToolkit REQUIRED) +else() + set(TORCHLEAN_LIBTORCH_CUDA 0) endif() get_target_property(torch_cxx_standard torch CXX_STANDARD) @@ -55,17 +60,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}) @@ -75,14 +88,17 @@ 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}" "${CMAKE_CURRENT_SOURCE_DIR}" ) -target_link_libraries(torchlean_libtorch PRIVATE ${TORCH_LIBRARIES} - CUDA::nvrtc CUDA::cuda_driver) +target_link_libraries(torchlean_libtorch PRIVATE ${TORCH_LIBRARIES}) +if(TORCHLEAN_LIBTORCH_CUDA) + target_link_libraries(torchlean_libtorch PRIVATE CUDA::nvrtc CUDA::cuda_driver) +endif() set_target_properties(torchlean_libtorch PROPERTIES POSITION_INDEPENDENT_CODE ON LIBRARY_OUTPUT_DIRECTORY "${TORCHLEAN_OUTPUT_DIR}" @@ -98,6 +114,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} @@ -108,5 +125,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 b400310f6..54d89fcb8 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 @@ -30,12 +32,40 @@ do not record gradients. LayerNorm continues to use the existing normalization p ## 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. See the [LibTorch build instructions](../../scripts/README.md#libtorch-cuda-build) for SDK selection, prerequisites, compiler controls, cache selection, and build/test commands. Generated custom operations also link the toolkit's NVRTC and CUDA driver interfaces. They compile CUDA source at runtime; users do not maintain separate handwritten kernels for those operations. +A CPU-only SDK links neither, and the bridge built from it has no implementation of those +operations; see the host device below. + +### Host device + +A CPU-only SDK builds the same backend without CUDA support (CMake sees no `torch_cuda` target, +compiles with `TORCHLEAN_LIBTORCH_CUDA=0`, and links neither NVRTC nor the CUDA driver), and the +bridge then runs every ATen 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, generated custom operations fail with an error at their native entry (their kernels are +compiled at run time for a CUDA device), 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 diff --git a/csrc/libtorch/torchlean.cpp b/csrc/libtorch/torchlean.cpp index 4aaa9d3fc..08d6c8db0 100644 --- a/csrc/libtorch/torchlean.cpp +++ b/csrc/libtorch/torchlean.cpp @@ -4,18 +4,27 @@ // 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 #include #include -#include -#include #include #include +#endif +#include +#include #include #include +#include #include #include #include @@ -49,6 +58,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; @@ -181,6 +223,8 @@ lean_obj_res download(b_lean_obj_arg object) { return out; } +#if TORCHLEAN_LIBTORCH_CUDA +// Runtime-compiled kernels: NVRTC and the CUDA driver API come with a CUDA-enabled SDK only. void check_driver(CUresult result) { if (result == CUDA_SUCCESS) return; const char* message = nullptr; @@ -263,6 +307,9 @@ std::shared_ptr compile_kernel(const std::string& source) { at::Tensor run_kernel(b_lean_obj_arg source, const std::vector& inputs, at::ScalarType scalar, uint64_t count, uint64_t stride = 1) { + // A generated kernel reads device pointers: the host has no implementation of it. + TORCH_CHECK(selection == Selection::cuda, + "LibTorch: a runtime-compiled kernel needs a CUDA device, but none is selected"); TORCH_CHECK(stride > 0 && stride <= INT64_MAX && count <= INT64_MAX / stride, "kernel: invalid scalar width or output extent"); TORCH_CHECK(count <= INT64_MAX, "kernel: output size exceeds signed tensor size range"); @@ -309,6 +356,15 @@ at::Tensor run_kernel(b_lean_obj_arg source, const std::vector& inpu static_cast(record[3])); return output; } +#else +// A CPU-only SDK has no NVRTC: a runtime-compiled kernel fails at its entry. +at::Tensor run_kernel(b_lean_obj_arg, const std::vector&, at::ScalarType, uint64_t, + uint64_t = 1) { + TORCH_CHECK(false, "LibTorch: a runtime-compiled kernel needs a CUDA device, but this SDK is " + "CPU-only"); + return at::Tensor(); +} +#endif at::Tensor upload_bytes(b_lean_obj_arg object, at::ScalarType dtype, at::ScalarType source_dtype) { const size_t bytes = lean_sarray_size(object); @@ -357,16 +413,29 @@ lean_obj_res download_bytes(b_lean_obj_arg object, at::ScalarType dtype) { 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. @@ -393,11 +462,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())); } @@ -408,10 +479,10 @@ 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() && + TORCH_CHECK(buffer->tensor.device().type() == device().type() && (buffer->tensor.scalar_type() == at::kFloat || buffer->tensor.scalar_type() == at::kDouble), - "LibTorch: expected a CUDA real buffer"); + "LibTorch: expected a real 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"); @@ -419,9 +490,9 @@ const at::Tensor& tensor(b_lean_obj_arg object) { } torchlean_cuda_buffer* owned(at::Tensor value) { - TORCH_CHECK(value.defined() && value.is_cuda() && + TORCH_CHECK(value.defined() && value.device().type() == device().type() && (value.scalar_type() == at::kFloat || value.scalar_type() == at::kDouble), - "LibTorch: expected a CUDA real result"); + "LibTorch: expected a real 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(); @@ -437,7 +508,8 @@ lean_obj_res box(at::Tensor value) { } lean_obj_res encoded(at::Tensor value, b_lean_obj_arg format, size_t width) { - TORCH_CHECK(value.is_cuda() && value.scalar_type() == at::kByte && width > 0 && + TORCH_CHECK(value.device().type() == device().type() && value.scalar_type() == at::kByte && + width > 0 && value.numel() % width == 0, "encoded buffer: invalid storage extent"); TORCH_CHECK(static_cast(value.numel()) / width <= UINT32_MAX, "encoded buffer: scalar count exceeds the tape ABI"); @@ -480,8 +552,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) \ @@ -506,10 +594,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) @@ -744,7 +837,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) { @@ -753,8 +846,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)); @@ -816,8 +909,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); }); } @@ -825,21 +922,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 a335fc2bf..d81db5310 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 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 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 9867a2a96..81e7aa7e3 100644 --- a/csrc/libtorch/unavailable.c +++ b/csrc/libtorch/unavailable.c @@ -31,12 +31,19 @@ LEAN_EXPORT lean_obj_res torchlean_kernel_run_bytes( return unavailable_io(); } -// 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 45e5bbb59..b89bff12d 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" @@ -37,9 +38,11 @@ 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 rejects GPU computation. Status and memory queries remain available. +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 rejects GPU computation. Status and memory queries remain +available. -/ /-- @@ -63,7 +66,8 @@ private def nativeCompilerJob (name : String) : SpawnM (Job FilePath) := Job.asy return compiler error s!"native compiler not found: {name}" -/-- 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/README.md b/scripts/README.md index 6ba0832fd..45eec26af 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -70,7 +70,10 @@ The CUDA build needs a Linux host with a CUDA-capable GPU and a full LibTorch SD Attention and its local VJP are composed in Lean from LibTorch numerical primitives; there is no fused attention selector. Rerun the GPU checks when changing SDKs. -`cuda=true` selects the complete LibTorch backend. The SDK root must contain `include/`, +`cuda=true` selects the complete LibTorch backend; the SDK it is built against decides the +device, and a CPU-only SDK runs it on the host (see the +[backend's host device notes](../csrc/libtorch/README.md#host-device)). The SDK root must contain +`include/`, `lib/`, and `share/cmake/Torch/TorchConfig.cmake`. A CUDA-enabled Python PyTorch installation can provide this root. Partial header snapshots alone cannot build the backend. @@ -94,8 +97,8 @@ scripts/lake.sh -Kcuda=true lint The sanitizer wrapper enables this mode automatically. GPU execution needs a supported visible device; SDK configuration may also probe visible devices. TorchLean builds one shared C++ adapter and has no independent -nvcc prerequisite. The SDK's own CMake package may require a matching CUDA development toolkit -and invoke its compiler during discovery. +nvcc prerequisite. A CUDA-enabled SDK's own CMake package may require a matching CUDA development +toolkit and invoke its compiler during discovery. Build controls: diff --git a/scripts/libtorch_build.py b/scripts/libtorch_build.py index 518db904f..a636b3dc8 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 a6d46f28579ad9a4bdeabd2cc2ca86b4fdf23d25 Mon Sep 17 00:00:00 2001 From: Nicolas Rouquette Date: Sat, 3 Oct 2026 18:25:29 -0700 Subject: [PATCH 2/5] 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 | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/csrc/libtorch/README.md b/csrc/libtorch/README.md index 54d89fcb8..e61d687e3 100644 --- a/csrc/libtorch/README.md +++ b/csrc/libtorch/README.md @@ -64,8 +64,10 @@ rather than falling back to the host. On the host, the native allocator counters read zero, `synchronize` and `emptyCache` return at once, `setMemoryFraction` and `setDevice` are rejected, generated custom operations fail with an error at their native entry (their kernels are compiled at run time for a CUDA device), 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. +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 From 5555080f4aef125e30557fb484a2616f5ca0d2b1 Mon Sep 17 00:00:00 2001 From: Nicolas Rouquette Date: Sun, 18 Oct 2026 15:58:15 -0700 Subject: [PATCH 3/5] fix(tests): require a CUDA device under TORCHLEAN_REQUIRE_CUDA=1 `nativeAvailable` covers the host now that the bridge can run there, so `requireNativeRuntime`, which the suite called for `TORCHLEAN_REQUIRE_CUDA=1`, let a run on a CPU-only SDK or under `TORCHLEAN_LIBTORCH_DEVICE=cpu` pass a check meant to say that a GPU ran it. `Buffer.requireCudaDevice` asks for the device kind as well, and the suite uses it for that variable. `TORCHLEAN_REQUIRE_NATIVE=1` keeps the bridge-only meaning for CPU-only SDK sessions, so that such a job fails on the unlinked default backend instead of passing with the device suite skipped. In the bridge, `require_cuda("")` is the entry check for an operation with no host implementation (device selection, the allocator memory fraction and runtime-compiled kernels use it), so it fails with a message naming the feature rather than with a CUDA call on a CPU tensor. The READMEs, the installation page and the execution-modes chapter say what each variable rejects, and quote the new messages. Configured binary arithmetic computes only in runtime-compiled kernels, which the tape reaches through pure entry points, where an error ends the process. The encoded upload therefore calls `require_cuda` too, so on the host such a computation fails through `IO` when its tensor is uploaded. The host branch of the CUDA suite checks both refusals and runs the custom-autograd checks that need no generated kernel: the native-precision recordings and the resident-buffer checks. --- .../Autograd/Engine/LibTorch/Buffer.lean | 25 ++++++++++++++--- NN/Tests/Runtime/Cuda/Elementwise.lean | 19 ++++++++++++- NN/Tests/Suite.lean | 5 ++++ csrc/libtorch/README.md | 15 ++++++++++- csrc/libtorch/torchlean.cpp | 27 ++++++++++++++----- .../Guide/Ch2_Frontend/ExecutionModes.lean | 15 ++++++----- home_page/installation/index.md | 3 ++- scripts/README.md | 7 ++--- 8 files changed, 92 insertions(+), 24 deletions(-) diff --git a/NN/Runtime/Autograd/Engine/LibTorch/Buffer.lean b/NN/Runtime/Autograd/Engine/LibTorch/Buffer.lean index d816fe937..804e08b51 100644 --- a/NN/Runtime/Autograd/Engine/LibTorch/Buffer.lean +++ b/NN/Runtime/Autograd/Engine/LibTorch/Buffer.lean @@ -78,18 +78,35 @@ private opaque runtimeStatusRaw (token : UInt32) : UInt32 | 1 => .nativeAvailable | _ => .nativeUnavailable -/-- Require the LibTorch bridge, with a device selected, 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 LibTorch build, but no usable CUDA device is visible; \ - TORCHLEAN_LIBTORCH_DEVICE=cpu selects the host instead" + "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 -/ diff --git a/NN/Tests/Runtime/Cuda/Elementwise.lean b/NN/Tests/Runtime/Cuda/Elementwise.lean index 475845d2b..df67e6495 100644 --- a/NN/Tests/Runtime/Cuda/Elementwise.lean +++ b/NN/Tests/Runtime/Cuda/Elementwise.lean @@ -736,7 +736,8 @@ private def checkCustomAutograd : IO Unit := do checkRecordingBuffers /-- On the host a generated operation has no implementation: it must fail through `IO`, naming -the requirement, before any CUDA call. -/ +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 @@ -748,6 +749,17 @@ private def checkCustomComputationsOnHost : IO Unit := do | 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. -/ @@ -760,6 +772,11 @@ private def checkCustomComputationsOnHost : IO Unit := do 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 := diff --git a/NN/Tests/Suite.lean b/NN/Tests/Suite.lean index 0cff42880..4b06d92b8 100644 --- a/NN/Tests/Suite.lean +++ b/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. diff --git a/csrc/libtorch/README.md b/csrc/libtorch/README.md index e61d687e3..1670aa2c9 100644 --- a/csrc/libtorch/README.md +++ b/csrc/libtorch/README.md @@ -53,7 +53,7 @@ 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 +TORCHLEAN_REQUIRE_NATIVE=1 scripts/lake.sh -Kcuda=true test ``` The device is selected once, at the bridge's first call, and `Runtime.Autograd.LibTorch.deviceKind` @@ -69,6 +69,19 @@ operations, gathers and lookups: ATen's CPU and CUDA transcendental functions (` `tanh`, ...) come from different libraries and differ at the ulp level, and reductions differ in order. The suite's tolerances hold on both. +Two environment variables make a run's expectation explicit. `TORCHLEAN_REQUIRE_NATIVE=1` fails +the suite unless the bridge is linked and selected a device, CUDA or host, so a CPU-only SDK job +cannot pass on the unlinked default backend. `TORCHLEAN_REQUIRE_CUDA=1` additionally requires +that device to be a CUDA device: it fails on the host, so a run that claims GPU coverage was one. +In the native code, everything that needs the CUDA headers compiles under +`#if TORCHLEAN_LIBTORCH_CUDA`, and an entry point with no host implementation calls +`require_cuda("")` before anything else, so it fails with a message naming the feature +rather than with a CUDA call on a CPU tensor. Kernels compiled at run time follow both rules: +NVRTC and the driver API sit under the macro, and `run_kernel` calls `require_cuda` at its entry. +An encoded buffer, the storage of configured binary arithmetic, is refused on the host where it is +created: only generated kernels compute on it, and the tape reaches them through pure entry points, +where an error would end the process instead of returning through `IO`. + ## Tested SDK versions Recorded CUDA-suite and focused attention validation used an A100 with pip torch 2.11.0+cu128. diff --git a/csrc/libtorch/torchlean.cpp b/csrc/libtorch/torchlean.cpp index 08d6c8db0..171b0ce87 100644 --- a/csrc/libtorch/torchlean.cpp +++ b/csrc/libtorch/torchlean.cpp @@ -91,6 +91,18 @@ Selection select_device() { return TORCHLEAN_LIBTORCH_CUDA ? Selection::unavailable : Selection::host; } +// An entry point with no host implementation calls this before anything else, so that it fails +// with a message naming the feature rather than with a CUDA call on a CPU tensor. Code that needs +// the CUDA headers to compile sits under TORCHLEAN_LIBTORCH_CUDA; code that merely needs the CUDA +// device selected calls this. +void require_cuda(const char* feature) { + TORCH_CHECK(selection != Selection::host, "LibTorch: ", feature, + " needs a CUDA device, but the bridge runs on the host (a CPU-only SDK, or " + "TORCHLEAN_LIBTORCH_DEVICE=cpu)"); + TORCH_CHECK(selection == Selection::cuda, "LibTorch: ", feature, + " needs a CUDA device, but none is selected"); +} + bool release_data(torchlean_cuda_buffer* buffer) { if (!buffer || !buffer->tensor.defined()) return false; const size_t size = buffer->size; @@ -308,8 +320,7 @@ std::shared_ptr compile_kernel(const std::string& source) { at::Tensor run_kernel(b_lean_obj_arg source, const std::vector& inputs, at::ScalarType scalar, uint64_t count, uint64_t stride = 1) { // A generated kernel reads device pointers: the host has no implementation of it. - TORCH_CHECK(selection == Selection::cuda, - "LibTorch: a runtime-compiled kernel needs a CUDA device, but none is selected"); + require_cuda("a runtime-compiled kernel"); TORCH_CHECK(stride > 0 && stride <= INT64_MAX && count <= INT64_MAX / stride, "kernel: invalid scalar width or output extent"); TORCH_CHECK(count <= INT64_MAX, "kernel: output size exceeds signed tensor size range"); @@ -360,8 +371,8 @@ at::Tensor run_kernel(b_lean_obj_arg source, const std::vector& inpu // A CPU-only SDK has no NVRTC: a runtime-compiled kernel fails at its entry. at::Tensor run_kernel(b_lean_obj_arg, const std::vector&, at::ScalarType, uint64_t, uint64_t = 1) { - TORCH_CHECK(false, "LibTorch: a runtime-compiled kernel needs a CUDA device, but this SDK is " - "CPU-only"); + require_cuda("a runtime-compiled kernel"); + TORCH_CHECK(false, "LibTorch: a runtime-compiled kernel needs a CUDA-enabled SDK"); return at::Tensor(); } #endif @@ -788,6 +799,9 @@ extern "C" LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_format(b_lean_obj_arg extern "C" LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_of_encoded_io( b_lean_obj_arg bytes, b_lean_obj_arg format, uint64_t width) { return io([&] { + // Only runtime-compiled kernels compute on encoded words, and the tape reaches them through + // pure entry points, where an error would end the process: refuse the upload instead. + require_cuda("configured binary arithmetic"); TORCH_CHECK(width && lean_sarray_size(bytes) % width == 0, "encoded buffer: truncated scalar word"); return torchlean::encoded(upload_bytes(bytes, at::kByte, at::kByte), format, width); @@ -846,7 +860,7 @@ 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(selection == Selection::cuda, "LibTorch: no CUDA device is selected"); + require_cuda("device selection"); 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"); @@ -922,8 +936,7 @@ 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"); + require_cuda("the allocator memory fraction"); #if TORCHLEAN_LIBTORCH_CUDA c10::cuda::CUDACachingAllocator::setMemoryFraction(fraction, torchlean::device().index()); #endif diff --git a/home_page/blueprint/TorchLeanBlueprint/Guide/Ch2_Frontend/ExecutionModes.lean b/home_page/blueprint/TorchLeanBlueprint/Guide/Ch2_Frontend/ExecutionModes.lean index fe7466224..d4cdf13d4 100644 --- a/home_page/blueprint/TorchLeanBlueprint/Guide/Ch2_Frontend/ExecutionModes.lean +++ b/home_page/blueprint/TorchLeanBlueprint/Guide/Ch2_Frontend/ExecutionModes.lean @@ -1174,16 +1174,17 @@ A CUDA build with no visible device: $ # Hide visible GPUs while keeping the native CUDA build to $ # isolate runtime availability. $ CUDA_VISIBLE_DEVICES="" scripts/lake.sh -K cuda=true exe torchlean quickstart_mlp --device cuda -error: torch eager session: CUDA was requested and -this is a CUDA build, but no usable CUDA device is visible +error: torch eager session: 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 ``` That last message distinguishes the two failure modes it could be reporting. A build without -LibTorch says so and tells you to rebuild with `-K cuda=true`; a real CUDA build with no -device says that instead. Both refuse. Notice also where the second message came from: the typed -graph rejection happened after the banner line had already printed, because it is raised when the -session opens rather than when the flags are parsed. Session validation checks programmatic -configurations as well as command-line requests. +LibTorch says so and tells you to rebuild with `-K cuda=true`; a LibTorch build with no device says +that instead, and names the way to run on the host. Both refuse. Notice also where the second +message came from: the typed graph rejection happened after the banner line had already printed, +because it is raised when the session opens rather than when the flags are parsed. Session +validation checks programmatic configurations as well as command-line requests. The rejection is a value in Lean, so this page can exhibit it without a terminal: diff --git a/home_page/installation/index.md b/home_page/installation/index.md index c6b42aad6..ae9d3b006 100644 --- a/home_page/installation/index.md +++ b/home_page/installation/index.md @@ -276,7 +276,8 @@ scripts/lake.sh exe verify --help For CUDA, rebuild with `scripts/lake.sh -Kcuda=true build`, then run the suite with `TORCHLEAN_REQUIRE_CUDA=1 scripts/lake.sh -Kcuda=true exe nn_tests_suite`. The environment -variable makes an unavailable CUDA runtime a failure instead of allowing CUDA tests to skip. +variable makes an unavailable CUDA runtime, or the bridge running on the host, a failure instead +of allowing CUDA tests to skip or to pass on the CPU. For a complete account of Lean axioms, executable checkers, CUDA and FFI code, external artifact producers, and floating-point assumptions, read diff --git a/scripts/README.md b/scripts/README.md index 45eec26af..22cb32248 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -93,9 +93,10 @@ scripts/lake.sh -Kcuda=true build NN.CI.All scripts/lake.sh -Kcuda=true lint ``` -`TORCHLEAN_REQUIRE_CUDA=1` rejects builds without LibTorch and missing CUDA devices. -The sanitizer wrapper enables this mode automatically. GPU execution needs a supported visible -device; SDK configuration may also probe +`TORCHLEAN_REQUIRE_CUDA=1` rejects builds without LibTorch, missing CUDA devices, and the bridge +running on the host; `TORCHLEAN_REQUIRE_NATIVE=1` rejects only the first two, for CPU-only SDK +sessions. The sanitizer wrapper enables `TORCHLEAN_REQUIRE_CUDA=1` automatically. GPU execution +needs a supported visible device; SDK configuration may also probe visible devices. TorchLean builds one shared C++ adapter and has no independent nvcc prerequisite. A CUDA-enabled SDK's own CMake package may require a matching CUDA development toolkit and invoke its compiler during discovery. From d7f5e85c8e1a7b579c47b603b2850da12d2ebec3 Mon Sep 17 00:00:00 2001 From: Nicolas Rouquette Date: Sat, 3 Oct 2026 12:01:30 -0700 Subject: [PATCH 4/5] feat(libtorch): report device identity per index `LibTorch.deviceInfo index` reads one device's name, compute capability, SM count, clocks, bus width, total memory, and the driver and runtime versions through the SDK's per-device cached properties; `currentDeviceInfo` reads the device `getDevice` selects. Querying by explicit index keeps the answer independent of the selected device, so switching devices cannot report a stale record. Without LibTorch, the queries fail through `IO` like the other controls. `NN/Tests/Runtime/Cuda/DeviceInfo.lean` checks every visible device under its own index, the current-device agreement, index rejection, and, in a fresh process selected by `TORCHLEAN_LIBTORCH_DEVICE_PROBE=switch`, that switching devices changes the reading. --- NN/Runtime/Autograd/Engine/LibTorch.lean | 1 + .../Autograd/Engine/LibTorch/DeviceInfo.lean | 137 ++++++++++++++++++ NN/Tests/Runtime/Cuda/DeviceInfo.lean | 111 ++++++++++++++ NN/Tests/Runtime/Cuda/Suite.lean | 1 + NN/Tests/Suite.lean | 5 + csrc/libtorch/README.md | 1 + csrc/libtorch/torchlean.cpp | 70 +++++++++ csrc/libtorch/unavailable.c | 9 ++ 8 files changed, 335 insertions(+) create mode 100644 NN/Runtime/Autograd/Engine/LibTorch/DeviceInfo.lean create mode 100644 NN/Tests/Runtime/Cuda/DeviceInfo.lean diff --git a/NN/Runtime/Autograd/Engine/LibTorch.lean b/NN/Runtime/Autograd/Engine/LibTorch.lean index ca48200d0..fe4f20300 100644 --- a/NN/Runtime/Autograd/Engine/LibTorch.lean +++ b/NN/Runtime/Autograd/Engine/LibTorch.lean @@ -14,6 +14,7 @@ public import NN.Runtime.Autograd.Engine.LibTorch.Float32Contract public import NN.Runtime.Autograd.Engine.LibTorch.KernelSpec public import NN.Runtime.Autograd.Engine.LibTorch.Kernels public import NN.Runtime.Autograd.Engine.LibTorch.Controls +public import NN.Runtime.Autograd.Engine.LibTorch.DeviceInfo public import NN.Runtime.Autograd.Engine.LibTorch.Ops public import NN.Runtime.Autograd.Engine.LibTorch.Shape public import NN.Runtime.Autograd.Engine.LibTorch.Tape diff --git a/NN/Runtime/Autograd/Engine/LibTorch/DeviceInfo.lean b/NN/Runtime/Autograd/Engine/LibTorch/DeviceInfo.lean new file mode 100644 index 000000000..7a68ff7e3 --- /dev/null +++ b/NN/Runtime/Autograd/Engine/LibTorch/DeviceInfo.lean @@ -0,0 +1,137 @@ +/- +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.Controls + +/-! +# LibTorch CUDA Device Identity + +Readback of the CUDA device the bridge runs on. A measurement that cannot name the device it was +taken on produces a number nobody can file beside another one. `Buffer.memory` answers how +much device memory there is; this module answers which device. + +Every query takes an explicit device index and reads that device's properties through the SDK's +own per-device cache, so the answer never depends on which device was selected earlier, and a +multi-GPU host never combines fields from two cards in one record. Without LibTorch linked, every +query fails through `IO` like the other runtime controls. + +The SDK's packaged GPU architectures are not reported: ATen's C++ API does not expose the list the +Python package reads back, and TorchLean compiles no device code of its own. +-/ + +@[expose] public section + +namespace Runtime +namespace Autograd +namespace LibTorch + +@[never_extract, extern "torchlean_libtorch_device_name"] +private opaque deviceNameRaw (index : UInt32) : IO String + +@[never_extract, extern "torchlean_libtorch_device_capability"] +private opaque deviceCapabilityRaw (index : UInt32) : IO UInt32 + +@[never_extract, extern "torchlean_libtorch_device_sm_count"] +private opaque deviceSmCountRaw (index : UInt32) : IO UInt32 + +@[never_extract, extern "torchlean_libtorch_device_clock_khz"] +private opaque deviceClockKhzRaw (index : UInt32) : IO UInt32 + +@[never_extract, extern "torchlean_libtorch_device_mem_clock_khz"] +private opaque deviceMemClockKhzRaw (index : UInt32) : IO UInt32 + +@[never_extract, extern "torchlean_libtorch_device_mem_bus_width"] +private opaque deviceMemBusWidthRaw (index : UInt32) : IO UInt32 + +@[never_extract, extern "torchlean_libtorch_device_total_bytes"] +private opaque deviceTotalBytesRaw (index : UInt32) : IO UInt64 + +@[never_extract, extern "torchlean_libtorch_driver_version"] +private opaque driverVersionRaw (token : UInt32) : IO UInt32 + +@[never_extract, extern "torchlean_libtorch_runtime_version"] +private opaque runtimeVersionRaw (token : UInt32) : IO UInt32 + +/-- +Identity of one CUDA device, as the linked SDK reports it. + +The device fields come from the SDK's cached `cudaDeviceProp` for that index and from +`cudaDeviceGetAttribute`; the two version fields from `cudaDriverGetVersion` and +`cudaRuntimeGetVersion`. `capability` is `major * 10 + minor`, so it reads the way an `sm_XY` +target is spelled. +-/ +structure DeviceInfo where + /-- The index this record was read for. -/ + index : UInt32 + /-- The device name, for example `NVIDIA RTX A4500`. -/ + name : String + /-- `major * 10 + minor`: `86` is `sm_86`, `120` is `sm_120`. -/ + capability : UInt32 + /-- Streaming multiprocessor count. -/ + smCount : UInt32 + /-- Core clock in kHz. -/ + clockKhz : UInt32 + /-- Memory clock in kHz. -/ + memClockKhz : UInt32 + /-- Global memory bus width in bits. -/ + memBusWidthBits : UInt32 + /-- Total global memory in bytes. -/ + totalBytes : UInt64 + /-- `cudaDriverGetVersion`, for example `13000` for 13.0. -/ + driverVersion : UInt32 + /-- `cudaRuntimeGetVersion` of the SDK's runtime, for example `12080` for 12.8. -/ + runtimeVersion : UInt32 + deriving Repr + +/-- +Read the properties of device `index`. + +Fails through `IO` when the index is outside the visible device count or LibTorch is not linked. +Reading by explicit index is what keeps the answer independent of the selected device. +-/ +@[no_expose] def deviceInfo (index : UInt32) : IO DeviceInfo := do + let name ← deviceNameRaw index + let capability ← deviceCapabilityRaw index + let smCount ← deviceSmCountRaw index + let clockKhz ← deviceClockKhzRaw index + let memClockKhz ← deviceMemClockKhzRaw index + let memBusWidthBits ← deviceMemBusWidthRaw index + let totalBytes ← deviceTotalBytesRaw index + let driverVersion ← driverVersionRaw 0 + let runtimeVersion ← runtimeVersionRaw 0 + pure { index, name, capability, smCount, clockKhz, memClockKhz, memBusWidthBits, totalBytes, + driverVersion, runtimeVersion } + +/-- The device the bridge currently selects, as `getDevice` reports it. -/ +@[no_expose] def currentDeviceInfo : IO DeviceInfo := do + deviceInfo (← getDevice) + +/-- A CUDA version integer (`12080`) as it is written (`12.8`). -/ +def versionString (v : UInt32) : String := + let n := v.toNat + s!"{n / 1000}.{(n % 1000) / 10}" + +/-- +Peak theoretical memory bandwidth in GB/s from the memory clock and bus width, so a row reporting +achieved GB/s can be read as a fraction of the roofline without looking the card up. Double data +rate is assumed. +-/ +def DeviceInfo.peakBandwidthGBs (d : DeviceInfo) : Float := + (Float.ofNat d.memClockKhz.toNat) * 1.0e3 * 2.0 * + (Float.ofNat d.memBusWidthBits.toNat / 8.0) / 1.0e9 + +/-- One-line device identity for a benchmark header. -/ +def DeviceInfo.format (d : DeviceInfo) : String := + s!"{d.name} (device {d.index}, sm_{d.capability}, {d.smCount} SMs, " ++ + s!"{(Float.ofNat d.totalBytes.toNat) / (1024.0 * 1024.0 * 1024.0)} GiB, " ++ + s!"~{d.peakBandwidthGBs} GB/s peak, " ++ + s!"driver {versionString d.driverVersion}, runtime {versionString d.runtimeVersion})" + +end LibTorch +end Autograd +end Runtime diff --git a/NN/Tests/Runtime/Cuda/DeviceInfo.lean b/NN/Tests/Runtime/Cuda/DeviceInfo.lean new file mode 100644 index 000000000..d71b2b898 --- /dev/null +++ b/NN/Tests/Runtime/Cuda/DeviceInfo.lean @@ -0,0 +1,111 @@ +/- +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.Buffer +public import NN.Runtime.Autograd.Engine.LibTorch.DeviceInfo + +/-! +# Device identity readback + +Checks that every visible device reports itself under its own index, that the current-device +reading agrees with `getDevice`, that an index at the visible count is rejected, and, in a fresh +process, that switching devices changes what `currentDeviceInfo` reports. The switching probe runs +in a child process because `setDevice` refuses while any buffer wrapper is live, and the suite +allocates buffers before the CUDA coverage tests run. Without LibTorch, the queries must fail like +every other runtime control. +-/ + +@[expose] public section + +namespace Tests +namespace Cuda +namespace DeviceInfo + +open Runtime.Autograd.LibTorch + +def expect (label : String) (ok : Bool) : IO Unit := do + unless ok do + throw <| IO.userError s!"device identity check failed: {label}" + +/-- Every visible device reports itself under its own index; the current device agrees with +`getDevice`; an index at the visible count is rejected. -/ +def runReadback : IO Unit := do + let count ← deviceCount + expect "a native-available build sees at least one device" (count > 0) + for i in [0:count.toNat] do + let info ← deviceInfo i.toUInt32 + expect s!"device {i} reports its own index" (info.index == i.toUInt32) + expect s!"device {i} has a name" (!info.name.isEmpty) + expect s!"device {i} reports a compute capability" (info.capability > 0) + expect s!"device {i} reports an SM count" (info.smCount > 0) + expect s!"device {i} reports its memory" (info.totalBytes > 0) + IO.println s!" {info.format}" + let current ← currentDeviceInfo + expect "current device agrees with getDevice" (current.index == (← getDevice)) + let rejected ← + try + discard (deviceInfo count) + pure false + catch _ => pure true + expect "an index at the visible count is rejected" rejected + +/-- +Switch through every visible device and read the current device back after each switch. + +Runs under `TORCHLEAN_LIBTORCH_DEVICE_PROBE=switch` in a fresh process, before any buffer exists. +-/ +def runSwitchProbe : IO Unit := do + Buffer.requireNativeRuntime + let count ← deviceCount + let original ← getDevice + for i in [0:count.toNat] do + setDevice i.toUInt32 + let info ← currentDeviceInfo + expect s!"after setDevice {i}, currentDeviceInfo reports {i}" (info.index == i.toUInt32) + let direct ← deviceInfo i.toUInt32 + expect s!"after setDevice {i}, the name agrees with the direct read" (info.name == direct.name) + setDevice original + IO.println s!" device switching: {count} device(s) report under their own index" + +/-- Without LibTorch, identity queries fail like every other control. -/ +def runUnavailable : IO Unit := do + let failed ← + try + discard (deviceInfo 0) + pure false + catch _ => pure true + expect "deviceInfo fails without LibTorch" failed + IO.println " device identity: queries are rejected without LibTorch" + +/-- Entry point called from `NN/Tests/Suite.lean` under every runtime status. -/ +def run : IO Unit := do + IO.println "== LibTorch device identity ==" + match Buffer.runtimeStatus with + | .notLinked => runUnavailable + | .nativeUnavailable => pure () + | .nativeAvailable => + runReadback + let self : System.FilePath := "/proc/self/exe" + if !(← self.pathExists) then + IO.println " skipped: the device switching probe requires Linux /proc/self/exe" + return + let result ← IO.Process.output { + cmd := self.toString + args := #[] + env := #[("TORCHLEAN_LIBTORCH_DEVICE_PROBE", some "switch"), + ("TORCHLEAN_REQUIRE_CUDA", some "1")] + } + if result.exitCode != 0 then + throw <| IO.userError <| + s!"device switching probe failed (exit {result.exitCode}):\n" ++ + s!"{result.stdout}\n{result.stderr}" + IO.print result.stdout + +end DeviceInfo +end Cuda +end Tests diff --git a/NN/Tests/Runtime/Cuda/Suite.lean b/NN/Tests/Runtime/Cuda/Suite.lean index b9cb35c53..beb98bab4 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.DeviceInfo public import NN.Tests.Runtime.Cuda.SelectiveScan public import NN.Tests.Runtime.Cuda.Fft public import NN.Tests.Runtime.Cuda.Stress diff --git a/NN/Tests/Suite.lean b/NN/Tests/Suite.lean index 4b06d92b8..383bd13a9 100644 --- a/NN/Tests/Suite.lean +++ b/NN/Tests/Suite.lean @@ -59,6 +59,10 @@ def run : IO Unit := do Runtime.Autograd.LibTorch.Buffer.requireCudaDevice if (← IO.getEnv "TORCHLEAN_REQUIRE_NATIVE") == some "1" then Runtime.Autograd.LibTorch.Buffer.requireNativeRuntime + -- Device switching needs a process with no live buffer wrappers, so it runs on its own. + if (← IO.getEnv "TORCHLEAN_LIBTORCH_DEVICE_PROBE") == some "switch" then + Tests.Cuda.DeviceInfo.runSwitchProbe + return -- 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. match ← IO.getEnv "TORCHLEAN_LIBTORCH_MEMORY_PROBE" with @@ -81,6 +85,7 @@ def run : IO Unit := do NN.Tests.Tensor.Storage.run NN.Tests.Tensor.Stack.run Tests.Floats.run + Tests.Cuda.DeviceInfo.run match Runtime.Autograd.LibTorch.Buffer.runtimeStatus with | .notLinked => IO.println " CUDA kernels: skipped (LibTorch not linked)" diff --git a/csrc/libtorch/README.md b/csrc/libtorch/README.md index 1670aa2c9..5fd774b59 100644 --- a/csrc/libtorch/README.md +++ b/csrc/libtorch/README.md @@ -227,6 +227,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/DeviceInfo.lean` | Per-index device identity readback, current-device agreement, index rejection, and a device-switching probe in a fresh process. | | `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 171b0ce87..84cb6f7d7 100644 --- a/csrc/libtorch/torchlean.cpp +++ b/csrc/libtorch/torchlean.cpp @@ -869,6 +869,76 @@ extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_set_device(uint32_t index }); } +// Device identity. Every query takes an explicit index and reads the SDK's per-device cached +// properties, so the answer never depends on the selected device and a multi-GPU host never +// combines fields from two cards in one record. +namespace { + +const cudaDeviceProp& device_properties(uint32_t index) { + TORCH_CHECK(index < static_cast(c10::cuda::device_count()), + "LibTorch: CUDA device index is outside the visible count"); + return *at::cuda::getDeviceProperties(static_cast(index)); +} + +// Clocks and bus width come from `cudaDeviceGetAttribute`: CUDA 13 removed `clockRate` and +// `memoryClockRate` from `cudaDeviceProp`, and the attribute enums exist in 12.x and 13.x. +uint32_t device_attribute(uint32_t index, cudaDeviceAttr attribute) { + device_properties(index); + int value = 0; + C10_CUDA_CHECK(cudaDeviceGetAttribute(&value, attribute, static_cast(index))); + return static_cast(value); +} + +} // namespace + +extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_device_name(uint32_t index) { + return io([&] { return lean_mk_string(device_properties(index).name); }); +} + +extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_device_capability(uint32_t index) { + return io([&] { + const auto& properties = device_properties(index); + return lean_box_uint32(static_cast(properties.major * 10 + properties.minor)); + }); +} + +extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_device_sm_count(uint32_t index) { + return io([&] { + return lean_box_uint32(static_cast(device_properties(index).multiProcessorCount)); + }); +} + +extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_device_total_bytes(uint32_t index) { + return io([&] { + return lean_box_uint64(static_cast(device_properties(index).totalGlobalMem)); + }); +} + +#define TORCHLEAN_DEVICE_ATTRIBUTE(NAME, ATTRIBUTE) \ + extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_device_##NAME(uint32_t index) { \ + return io([&] { return lean_box_uint32(device_attribute(index, ATTRIBUTE)); }); \ + } +TORCHLEAN_DEVICE_ATTRIBUTE(clock_khz, cudaDevAttrClockRate) +TORCHLEAN_DEVICE_ATTRIBUTE(mem_clock_khz, cudaDevAttrMemoryClockRate) +TORCHLEAN_DEVICE_ATTRIBUTE(mem_bus_width, cudaDevAttrGlobalMemoryBusWidth) +#undef TORCHLEAN_DEVICE_ATTRIBUTE + +extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_driver_version(uint32_t) { + return io([] { + int version = 0; + C10_CUDA_CHECK(cudaDriverGetVersion(&version)); + return lean_box_uint32(static_cast(version)); + }); +} + +extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_runtime_version(uint32_t) { + return io([] { + int version = 0; + C10_CUDA_CHECK(cudaRuntimeGetVersion(&version)); + return lean_box_uint32(static_cast(version)); + }); +} + extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_get_setting(uint32_t setting) { return io([&] { auto& context = at::globalContext(); diff --git a/csrc/libtorch/unavailable.c b/csrc/libtorch/unavailable.c index 81e7aa7e3..c86f26534 100644 --- a/csrc/libtorch/unavailable.c +++ b/csrc/libtorch/unavailable.c @@ -95,6 +95,15 @@ ZERO(uint64_t, torchlean_cuda_allocator_device_total_bytes) LEAN_EXPORT lean_obj_res NAME(__VA_ARGS__) { return unavailable_io(); } UNAVAILABLE_IO(torchlean_libtorch_set_device, uint32_t index) +UNAVAILABLE_IO(torchlean_libtorch_device_name, uint32_t index) +UNAVAILABLE_IO(torchlean_libtorch_device_capability, uint32_t index) +UNAVAILABLE_IO(torchlean_libtorch_device_sm_count, uint32_t index) +UNAVAILABLE_IO(torchlean_libtorch_device_clock_khz, uint32_t index) +UNAVAILABLE_IO(torchlean_libtorch_device_mem_clock_khz, uint32_t index) +UNAVAILABLE_IO(torchlean_libtorch_device_mem_bus_width, uint32_t index) +UNAVAILABLE_IO(torchlean_libtorch_device_total_bytes, uint32_t index) +UNAVAILABLE_IO(torchlean_libtorch_driver_version, uint32_t token) +UNAVAILABLE_IO(torchlean_libtorch_runtime_version, uint32_t token) UNAVAILABLE_IO(torchlean_libtorch_set_setting, uint32_t setting, uint32_t enabled) UNAVAILABLE_IO(torchlean_libtorch_set_memory_fraction, double fraction) UNAVAILABLE_IO(torchlean_libtorch_synchronize, uint32_t token) From cfbcb9a57675b6a1a19f9b1a8b6b7628a4027445 Mon Sep 17 00:00:00 2001 From: Nicolas Rouquette Date: Sun, 18 Oct 2026 16:29:41 -0700 Subject: [PATCH 5/5] fix(libtorch): handle CPU-only SDKs and the host in the device identity queries The identity exports read `cudaDeviceProp` and `cudaDeviceGetAttribute`, which a CPU-only SDK does not provide, and the test read `nativeAvailable` as "a CUDA device is present", which the host bridge no longer implies. The exports now compile under `TORCHLEAN_LIBTORCH_CUDA`; a CPU-only SDK gets exports that fail through `IO`, as an index outside the visible count does on a CUDA-enabled SDK. A CUDA-enabled SDK still answers while the bridge runs on the host, since the queries describe the visible devices rather than the selected one. `NN/Tests/Runtime/Cuda/DeviceInfo.lean` follows the device the bridge selected: on a CUDA device it runs the readback and the switching probe, which now requires a CUDA device; on the host with a CUDA-enabled SDK it runs the readback alone, because the host rejects `setDevice`; on a CPU-only SDK it checks that every query fails. The module docstring keeps the SDK's packaged architectures, a property of the SDK build, apart from the kernels TorchLean compiles at run time: NVRTC compiles the generated operations of `NN.Kernel` for the compute capability of the device they run on. --- .../Autograd/Engine/LibTorch/DeviceInfo.lean | 18 +++++++--- NN/Tests/Runtime/Cuda/DeviceInfo.lean | 35 +++++++++++++------ csrc/libtorch/README.md | 2 +- csrc/libtorch/torchlean.cpp | 26 +++++++++++++- 4 files changed, 64 insertions(+), 17 deletions(-) diff --git a/NN/Runtime/Autograd/Engine/LibTorch/DeviceInfo.lean b/NN/Runtime/Autograd/Engine/LibTorch/DeviceInfo.lean index 7a68ff7e3..9a5659928 100644 --- a/NN/Runtime/Autograd/Engine/LibTorch/DeviceInfo.lean +++ b/NN/Runtime/Autograd/Engine/LibTorch/DeviceInfo.lean @@ -17,11 +17,19 @@ much device memory there is; this module answers which device. Every query takes an explicit device index and reads that device's properties through the SDK's own per-device cache, so the answer never depends on which device was selected earlier, and a -multi-GPU host never combines fields from two cards in one record. Without LibTorch linked, every -query fails through `IO` like the other runtime controls. - -The SDK's packaged GPU architectures are not reported: ATen's C++ API does not expose the list the -Python package reads back, and TorchLean compiles no device code of its own. +multi-GPU host never combines fields from two cards in one record. The queries describe the +visible CUDA devices, not the one the bridge selected: a CUDA-enabled SDK answers them even while +the bridge runs on the host (`TORCHLEAN_LIBTORCH_DEVICE=cpu`). Without LibTorch linked, and on a +CPU-only SDK, every query fails through `IO` like the other runtime controls. + +Two things this record does not say, which are separate questions. The SDK's packaged GPU +architectures, the `sm_XY` targets its own kernels were compiled for: ATen's C++ API does not +expose the list the Python package reads back, and it is a property of the SDK build rather than +of the device, so whether the packaged kernels cover a device is settled by comparing +`capability` against that list by other means. And the kernels TorchLean compiles at run time, +the generated operations of `NN.Kernel`: NVRTC compiles those for the compute capability of the +device they run on, whatever the SDK was packaged for, so their coverage is a question for NVRTC, +not for this readback. -/ @[expose] public section diff --git a/NN/Tests/Runtime/Cuda/DeviceInfo.lean b/NN/Tests/Runtime/Cuda/DeviceInfo.lean index d71b2b898..e0cfc08c4 100644 --- a/NN/Tests/Runtime/Cuda/DeviceInfo.lean +++ b/NN/Tests/Runtime/Cuda/DeviceInfo.lean @@ -16,8 +16,12 @@ Checks that every visible device reports itself under its own index, that the cu reading agrees with `getDevice`, that an index at the visible count is rejected, and, in a fresh process, that switching devices changes what `currentDeviceInfo` reports. The switching probe runs in a child process because `setDevice` refuses while any buffer wrapper is live, and the suite -allocates buffers before the CUDA coverage tests run. Without LibTorch, the queries must fail like -every other runtime control. +allocates buffers before the CUDA coverage tests run. + +`nativeAvailable` does not mean a CUDA device: the bridge may run on the host. On a CUDA-enabled +SDK whose bridge runs on the host the readback still answers for every visible device, but +`setDevice` is rejected there, so the switching probe does not run. On a CPU-only SDK no device is +visible and every query must fail, as it must without LibTorch. -/ @[expose] public section @@ -36,7 +40,7 @@ def expect (label : String) (ok : Bool) : IO Unit := do `getDevice`; an index at the visible count is rejected. -/ def runReadback : IO Unit := do let count ← deviceCount - expect "a native-available build sees at least one device" (count > 0) + expect "a CUDA-enabled SDK sees at least one device" (count > 0) for i in [0:count.toNat] do let info ← deviceInfo i.toUInt32 expect s!"device {i} reports its own index" (info.index == i.toUInt32) @@ -57,10 +61,11 @@ def runReadback : IO Unit := do /-- Switch through every visible device and read the current device back after each switch. -Runs under `TORCHLEAN_LIBTORCH_DEVICE_PROBE=switch` in a fresh process, before any buffer exists. +Runs under `TORCHLEAN_LIBTORCH_DEVICE_PROBE=switch` in a fresh process, before any buffer exists, +and only with the bridge on a CUDA device: the host rejects `setDevice`. -/ def runSwitchProbe : IO Unit := do - Buffer.requireNativeRuntime + Buffer.requireCudaDevice let count ← deviceCount let original ← getDevice for i in [0:count.toNat] do @@ -72,23 +77,33 @@ def runSwitchProbe : IO Unit := do setDevice original IO.println s!" device switching: {count} device(s) report under their own index" -/-- Without LibTorch, identity queries fail like every other control. -/ -def runUnavailable : IO Unit := do +/-- Without LibTorch, and on a CPU-only SDK, identity queries fail like every other control. -/ +def runUnavailable (why : String) : IO Unit := do let failed ← try discard (deviceInfo 0) pure false catch _ => pure true - expect "deviceInfo fails without LibTorch" failed - IO.println " device identity: queries are rejected without LibTorch" + expect s!"deviceInfo fails {why}" failed + IO.println s!" device identity: queries are rejected {why}" /-- Entry point called from `NN/Tests/Suite.lean` under every runtime status. -/ def run : IO Unit := do IO.println "== LibTorch device identity ==" match Buffer.runtimeStatus with - | .notLinked => runUnavailable + | .notLinked => runUnavailable "without LibTorch" | .nativeUnavailable => pure () | .nativeAvailable => + if (← deviceKind) == .host then + if (← deviceCount) == 0 then + -- A CPU-only SDK: there is no CUDA runtime to ask. + runUnavailable "on a CPU-only SDK" + return + -- A CUDA-enabled SDK whose bridge runs on the host still describes its visible devices, + -- but rejects `setDevice`, so the switching probe cannot run. + runReadback + IO.println " skipped: the device switching probe needs the bridge on a CUDA device" + return runReadback let self : System.FilePath := "/proc/self/exe" if !(← self.pathExists) then diff --git a/csrc/libtorch/README.md b/csrc/libtorch/README.md index 5fd774b59..a5f6c2bb5 100644 --- a/csrc/libtorch/README.md +++ b/csrc/libtorch/README.md @@ -227,7 +227,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/DeviceInfo.lean` | Per-index device identity readback, current-device agreement, index rejection, and a device-switching probe in a fresh process. | +| `NN/Tests/Runtime/Cuda/DeviceInfo.lean` | Per-index device identity readback, current-device agreement, index rejection, and a device-switching probe in a fresh process; with the bridge on the host, the readback alone on a CUDA-enabled SDK and the rejection of every query on a CPU-only SDK. | | `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 84cb6f7d7..9365fded4 100644 --- a/csrc/libtorch/torchlean.cpp +++ b/csrc/libtorch/torchlean.cpp @@ -869,9 +869,11 @@ extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_set_device(uint32_t index }); } +#if TORCHLEAN_LIBTORCH_CUDA // Device identity. Every query takes an explicit index and reads the SDK's per-device cached // properties, so the answer never depends on the selected device and a multi-GPU host never -// combines fields from two cards in one record. +// combines fields from two cards in one record. A CUDA-enabled SDK answers these even while the +// bridge runs on the host: they describe the visible devices, not the selected one. namespace { const cudaDeviceProp& device_properties(uint32_t index) { @@ -938,6 +940,28 @@ extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_runtime_version(uint32_t) return lean_box_uint32(static_cast(version)); }); } +#else +// A CPU-only SDK has no CUDA device to describe and no driver to version: every identity query +// fails, as an index outside the visible count does on a CUDA-enabled SDK. +#define TORCHLEAN_NO_DEVICE_IDENTITY(NAME) \ + extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_##NAME(uint32_t) { \ + return io([]() -> lean_obj_res { \ + TORCH_CHECK(false, "LibTorch: this SDK has no CUDA support, so there is no CUDA device " \ + "identity to read"); \ + return lean_box(0); \ + }); \ + } +TORCHLEAN_NO_DEVICE_IDENTITY(device_name) +TORCHLEAN_NO_DEVICE_IDENTITY(device_capability) +TORCHLEAN_NO_DEVICE_IDENTITY(device_sm_count) +TORCHLEAN_NO_DEVICE_IDENTITY(device_total_bytes) +TORCHLEAN_NO_DEVICE_IDENTITY(device_clock_khz) +TORCHLEAN_NO_DEVICE_IDENTITY(device_mem_clock_khz) +TORCHLEAN_NO_DEVICE_IDENTITY(device_mem_bus_width) +TORCHLEAN_NO_DEVICE_IDENTITY(driver_version) +TORCHLEAN_NO_DEVICE_IDENTITY(runtime_version) +#undef TORCHLEAN_NO_DEVICE_IDENTITY +#endif extern "C" LEAN_EXPORT lean_obj_res torchlean_libtorch_get_setting(uint32_t setting) { return io([&] {