Add Warp support to the compiled runtime’s ONNX and PT2 backends

.github/workflows/ci.yml

@@ -57,6 +57,15 @@ jobs:
       - name: Install package and test dependencies
         run: python -m pip install -e ".[dev,test,warp-cu12]" cmake ninja
 
+      - name: Select GPU ONNX Runtime
+        run: |
+          # The base CPU dependency and Warp GPU extra share package files.
+          # Keep the CUDA-compatible version resolved by the Warp extra.
+          ort_gpu_version=$(python -c "from importlib.metadata import version; print(version('onnxruntime-gpu'))")
+          python -m pip uninstall -y onnxruntime
+          python -m pip install --force-reinstall --no-deps "onnxruntime-gpu==$ort_gpu_version"
+          python -c "import onnxruntime as ort; assert 'CUDAExecutionProvider' in ort.get_available_providers(), ort.get_available_providers()"
+
       - name: Build Warp runtime
         run: leapp-build-warp-runtime
 

leapp/leapp_graph/custom_operator_registry/warp_operator/__init__.py

@@ -120,13 +120,19 @@ def _register() -> None:
         and hasattr(getattr(torch.ops, NAMESPACE), OP_NAME)
     ):
         _LIB.define(_SCHEMA)
-        torch.library.register_fake(QUALIFIED_NAME, fx_impl.warp_runner_fake, lib=_LIB)
-        _LIB.impl(OP_NAME, fx_impl.warp_runner_eager, "CompositeExplicitAutograd")
 
         _get_logger().debug(
             f"Registered custom op {QUALIFIED_NAME} with schema: {_SCHEMA}"
         )
 
+    # The native adapter may register the schema before Python is imported.
+    if not torch._C._dispatch_has_kernel_for_dispatch_key(QUALIFIED_NAME, "Meta"):
+        torch.library.register_fake(QUALIFIED_NAME, fx_impl.warp_runner_fake, lib=_LIB)
+    if not torch._C._dispatch_has_kernel_for_dispatch_key(
+        QUALIFIED_NAME, "CompositeExplicitAutograd"
+    ):
+        _LIB.impl(OP_NAME, fx_impl.warp_runner_eager, "CompositeExplicitAutograd")
+
     if wp is not None:
         onnx_impl.register_onnx_lowering()
     _register_export_hooks()

leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/CMakeLists.txt

@@ -14,22 +14,44 @@ endforeach()
 option(LEAPP_WARP_BUILD_ONNX "Build ONNX Runtime custom op" ON)
 option(LEAPP_WARP_BUILD_TORCH "Build PyTorch dispatcher extension" ON)
 
-find_package(Python3 REQUIRED COMPONENTS Interpreter)
+include(GNUInstallDirs)
+if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND NOT DEFINED CMAKE_INSTALL_RPATH)
+    set(CMAKE_INSTALL_RPATH "$ORIGIN")
+endif()
+
+if(NOT DEFINED CUDAToolkit_ROOT AND DEFINED ENV{CUDAToolkit_ROOT})
+    set(CUDAToolkit_ROOT "$ENV{CUDAToolkit_ROOT}")
+endif()
+if(CUDAToolkit_ROOT AND NOT EXISTS "${CUDAToolkit_ROOT}/bin/nvcc${CMAKE_EXECUTABLE_SUFFIX}")
+    message(FATAL_ERROR "CUDA compiler not found under CUDAToolkit_ROOT: ${CUDAToolkit_ROOT}")
+endif()
 find_package(CUDAToolkit REQUIRED)
 find_package(Threads REQUIRED)
+message(STATUS "Warp CUDA toolkit: ${CUDAToolkit_LIBRARY_ROOT} (${CUDAToolkit_VERSION})")
 
-execute_process(
-    COMMAND "${Python3_EXECUTABLE}" -c "import os, warp; print(os.path.dirname(warp.__file__))"
-    OUTPUT_VARIABLE WARP_PACKAGE_DIR
-    OUTPUT_STRIP_TRAILING_WHITESPACE
-    COMMAND_ERROR_IS_FATAL ANY
-)
+function(leapp_warp_python OUTPUT CODE)
+    find_package(Python3 REQUIRED COMPONENTS Interpreter)
+    execute_process(
+        COMMAND "${Python3_EXECUTABLE}" -c "${CODE}"
+        OUTPUT_VARIABLE VALUE
+        OUTPUT_STRIP_TRAILING_WHITESPACE
+        COMMAND_ERROR_IS_FATAL ANY
+    )
+    set(${OUTPUT} "${VALUE}" PARENT_SCOPE)
+endfunction()
 
-set(WARP_NATIVE_DIR "${WARP_PACKAGE_DIR}/native")
-set(WARP_BIN_DIR "${WARP_PACKAGE_DIR}/bin")
+set(WARP_ROOT "$ENV{WARP_ROOT}" CACHE PATH "Warp SDK with native/ and bin/")
+if(NOT WARP_ROOT)
+    leapp_warp_python(WARP_ROOT "import os, warp; print(os.path.dirname(warp.__file__))")
+endif()
+set(WARP_NATIVE_DIR "${WARP_ROOT}/native")
+set(WARP_BIN_DIR "${WARP_ROOT}/bin")
 if(WIN32)
     set(WARP_SHARED_LIBRARY "${WARP_BIN_DIR}/warp.dll")
-    set(WARP_IMPORT_LIBRARY "${WARP_BIN_DIR}/warp.lib" CACHE FILEPATH "Warp import library")
+    set(WARP_IMPORT_LIBRARY "$ENV{LEAPP_WARP_IMPORT_LIBRARY}" CACHE FILEPATH "Warp import library")
+    if(NOT WARP_IMPORT_LIBRARY)
+        set(WARP_IMPORT_LIBRARY "${WARP_BIN_DIR}/warp.lib")
+    endif()
 elseif(APPLE)
     set(WARP_SHARED_LIBRARY "${WARP_BIN_DIR}/libwarp.dylib")
 else()
@@ -43,17 +65,22 @@ if(NOT EXISTS "${WARP_NATIVE_DIR}/apic.h")
     message(FATAL_ERROR "Warp APIC header not found: ${WARP_NATIVE_DIR}/apic.h. Install a Warp build with APIC support (Warp 1.13+ expected).")
 endif()
 
-add_library(warp_native SHARED IMPORTED)
-set_target_properties(warp_native PROPERTIES
-    IMPORTED_LOCATION "${WARP_SHARED_LIBRARY}"
-)
-if(WIN32)
-    if(NOT EXISTS "${WARP_IMPORT_LIBRARY}")
-        message(FATAL_ERROR "Warp import library not found: ${WARP_IMPORT_LIBRARY}. Set WARP_IMPORT_LIBRARY to a valid warp.lib.")
+message(STATUS "Warp native library: ${WARP_SHARED_LIBRARY}")
+if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+    # Warp's wheel ships warp.so without a SONAME or the usual lib prefix.
+    # Exact-filename linking avoids recording the absolute SDK path in DT_NEEDED.
+    add_library(warp_native INTERFACE)
+    target_link_directories(warp_native INTERFACE "${WARP_BIN_DIR}")
+    target_link_libraries(warp_native INTERFACE "-l:warp.so")
+else()
+    add_library(warp_native SHARED IMPORTED)
+    set_target_properties(warp_native PROPERTIES IMPORTED_LOCATION "${WARP_SHARED_LIBRARY}")
+    if(WIN32)
+        if(NOT EXISTS "${WARP_IMPORT_LIBRARY}")
+            message(FATAL_ERROR "Warp import library not found: ${WARP_IMPORT_LIBRARY}. Set WARP_IMPORT_LIBRARY to a valid warp.lib.")
+        endif()
+        set_target_properties(warp_native PROPERTIES IMPORTED_IMPLIB "${WARP_IMPORT_LIBRARY}")
     endif()
-    set_target_properties(warp_native PROPERTIES
-        IMPORTED_IMPLIB "${WARP_IMPORT_LIBRARY}"
-    )
 endif()
 
 set(LEAPP_WARP_CORE_SOURCES
@@ -70,73 +97,62 @@ target_include_directories(leapp_warp_runtime_core PUBLIC
     "${WARP_NATIVE_DIR}"
 )
 target_link_libraries(leapp_warp_runtime_core PUBLIC warp_native CUDA::cudart Threads::Threads)
+# Use CUDA vector types when Warp headers are included with the CUDA SDK.
+target_compile_definitions(leapp_warp_runtime_core PRIVATE WP_ENABLE_CUDA=1)
 set_target_properties(leapp_warp_runtime_core PROPERTIES POSITION_INDEPENDENT_CODE ON)
 
 if(LEAPP_WARP_BUILD_ONNX)
-    execute_process(
-        COMMAND "${Python3_EXECUTABLE}" -c "import os, onnxruntime; print(os.path.dirname(onnxruntime.__file__))"
-        OUTPUT_VARIABLE ONNXRUNTIME_PACKAGE_DIR
-        OUTPUT_STRIP_TRAILING_WHITESPACE
-        COMMAND_ERROR_IS_FATAL ANY
-    )
-    execute_process(
-        COMMAND "${Python3_EXECUTABLE}" -c "import onnxruntime; print(onnxruntime.__version__)"
-        OUTPUT_VARIABLE ONNXRUNTIME_VERSION
-        OUTPUT_STRIP_TRAILING_WHITESPACE
-        COMMAND_ERROR_IS_FATAL ANY
-    )
-
-    set(ORT_CAPI_DIR "${ONNXRUNTIME_PACKAGE_DIR}/capi")
-    if(WIN32)
-        set(ORT_SHARED_LIBRARY "${ORT_CAPI_DIR}/onnxruntime.dll")
-    elseif(APPLE)
-        set(ORT_SHARED_LIBRARY "${ORT_CAPI_DIR}/libonnxruntime.${ONNXRUNTIME_VERSION}.dylib")
-    else()
-        set(ORT_SHARED_LIBRARY "${ORT_CAPI_DIR}/libonnxruntime.so.${ONNXRUNTIME_VERSION}")
-    endif()
-    if(NOT EXISTS "${ORT_SHARED_LIBRARY}")
-        message(FATAL_ERROR "ONNX Runtime shared library not found: ${ORT_SHARED_LIBRARY}")
-    endif()
-
-    set(ORT_INCLUDE_DIR "${CMAKE_BINARY_DIR}/onnxruntime_headers")
-    file(MAKE_DIRECTORY "${ORT_INCLUDE_DIR}")
-    foreach(ORT_HEADER IN ITEMS
-        onnxruntime_c_api.h
-        onnxruntime_ep_c_api.h
-        onnxruntime_error_code.h
-    )
-        set(ORT_HEADER_PATH "${ORT_INCLUDE_DIR}/${ORT_HEADER}")
-        if(NOT EXISTS "${ORT_HEADER_PATH}")
-            file(DOWNLOAD
-                "https://raw.githubusercontent.com/microsoft/onnxruntime/v${ONNXRUNTIME_VERSION}/include/onnxruntime/core/session/${ORT_HEADER}"
-                "${ORT_HEADER_PATH}"
-                STATUS ORT_HEADER_DOWNLOAD_STATUS
-                TLS_VERIFY ON
-            )
-            list(GET ORT_HEADER_DOWNLOAD_STATUS 0 ORT_HEADER_DOWNLOAD_CODE)
-            list(GET ORT_HEADER_DOWNLOAD_STATUS 1 ORT_HEADER_DOWNLOAD_MESSAGE)
-            if(NOT ORT_HEADER_DOWNLOAD_CODE EQUAL 0)
-                # Split out of onnxruntime_c_api.h in 1.29; older releases 404.
-                if(ORT_HEADER STREQUAL "onnxruntime_error_code.h")
-                    file(REMOVE "${ORT_HEADER_PATH}")
-                    message(STATUS
-                        "Optional ${ORT_HEADER} is not available for "
-                        "ONNX Runtime ${ONNXRUNTIME_VERSION}; skipping.")
-                else()
-                    message(FATAL_ERROR "Failed to download ${ORT_HEADER}: ${ORT_HEADER_DOWNLOAD_MESSAGE}")
+    set(ORT_INCLUDE_DIR "$ENV{ORT_INCLUDE_DIR}" CACHE PATH "ONNX Runtime C API headers")
+    set(LEAPP_WARP_ONNXRUNTIME_VERSION "" CACHE STRING "Target ONNX Runtime header version (defaults to the Python package)")
+    if(NOT ORT_INCLUDE_DIR)
+        set(ORT_HEADER_VERSION "${LEAPP_WARP_ONNXRUNTIME_VERSION}")
+        if(NOT ORT_HEADER_VERSION)
+            leapp_warp_python(ORT_HEADER_VERSION "import onnxruntime; print(onnxruntime.__version__)")
+        endif()
+        set(ORT_INCLUDE_DIR "${CMAKE_BINARY_DIR}/onnxruntime_headers/${ORT_HEADER_VERSION}")
+        file(MAKE_DIRECTORY "${ORT_INCLUDE_DIR}")
+        foreach(ORT_HEADER IN ITEMS
+            onnxruntime_c_api.h
+            onnxruntime_ep_c_api.h
+            onnxruntime_error_code.h
+        )
+            set(ORT_HEADER_PATH "${ORT_INCLUDE_DIR}/${ORT_HEADER}")
+            if(NOT EXISTS "${ORT_HEADER_PATH}")
+                file(DOWNLOAD
+                    "https://raw.githubusercontent.com/microsoft/onnxruntime/v${ORT_HEADER_VERSION}/include/onnxruntime/core/session/${ORT_HEADER}"
+                    "${ORT_HEADER_PATH}"
+                    STATUS ORT_HEADER_DOWNLOAD_STATUS
+                    TLS_VERIFY ON
+                )
+                list(GET ORT_HEADER_DOWNLOAD_STATUS 0 ORT_HEADER_DOWNLOAD_CODE)
+                list(GET ORT_HEADER_DOWNLOAD_STATUS 1 ORT_HEADER_DOWNLOAD_MESSAGE)
+                if(NOT ORT_HEADER_DOWNLOAD_CODE EQUAL 0)
+                    # Split out of onnxruntime_c_api.h in 1.29; older releases 404.
+                    if(ORT_HEADER STREQUAL "onnxruntime_error_code.h")
+                        file(REMOVE "${ORT_HEADER_PATH}")
+                        message(STATUS
+                            "Optional ${ORT_HEADER} is not available for "
+                            "ONNX Runtime ${ORT_HEADER_VERSION}; skipping.")
+                    else()
+                        message(FATAL_ERROR "Failed to download ${ORT_HEADER}: ${ORT_HEADER_DOWNLOAD_MESSAGE}")
+                    endif()
                 endif()
             endif()
+        endforeach()
+    endif()
+    foreach(ORT_HEADER IN ITEMS onnxruntime_c_api.h onnxruntime_ep_c_api.h)
+        if(NOT EXISTS "${ORT_INCLUDE_DIR}/${ORT_HEADER}")
+            message(FATAL_ERROR "ONNX Runtime header not found: ${ORT_INCLUDE_DIR}/${ORT_HEADER}. Set ORT_INCLUDE_DIR to the target SDK headers.")
         endif()
     endforeach()
+    message(STATUS "Warp ONNX adapter headers: ${ORT_INCLUDE_DIR}")
 
     add_library(leapp_wrp_onnx_custom_op SHARED onnx/ort_wrp_runner_op.cc)
     target_include_directories(leapp_wrp_onnx_custom_op PRIVATE "${ORT_INCLUDE_DIR}")
     target_link_libraries(leapp_wrp_onnx_custom_op PRIVATE leapp_warp_runtime_core)
     if(NOT WIN32)
-        set_target_properties(leapp_wrp_onnx_custom_op PROPERTIES
-            BUILD_RPATH "${WARP_BIN_DIR};${CUDAToolkit_LIBRARY_DIR}"
-            INSTALL_RPATH "${WARP_BIN_DIR};${CUDAToolkit_LIBRARY_DIR}"
-        )
+        set_property(TARGET leapp_wrp_onnx_custom_op APPEND PROPERTY
+            BUILD_RPATH "${WARP_BIN_DIR};${CUDAToolkit_LIBRARY_DIR}")
     endif()
 endif()
 
@@ -147,21 +163,32 @@ if(LEAPP_WARP_BUILD_TORCH)
     if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
         set(CMAKE_CUDA_ARCHITECTURES native CACHE STRING "CUDA architectures for Torch-enabled builds")
     endif()
-    execute_process(
-        COMMAND "${Python3_EXECUTABLE}" -c "import torch; print(torch.utils.cmake_prefix_path)"
-        OUTPUT_VARIABLE TORCH_CMAKE_PREFIX_PATH
-        OUTPUT_STRIP_TRAILING_WHITESPACE
-        COMMAND_ERROR_IS_FATAL ANY
-    )
-    list(APPEND CMAKE_PREFIX_PATH "${TORCH_CMAKE_PREFIX_PATH}")
-    find_package(Torch REQUIRED)
+    if(NOT Torch_DIR AND DEFINED ENV{Torch_DIR})
+        set(Torch_DIR "$ENV{Torch_DIR}")
+    endif()
+    if(Torch_DIR AND NOT EXISTS "${Torch_DIR}/TorchConfig.cmake")
+        message(FATAL_ERROR "TorchConfig.cmake not found under Torch_DIR: ${Torch_DIR}")
+    endif()
+    find_package(Torch QUIET CONFIG)
+    if(NOT Torch_FOUND)
+        leapp_warp_python(TORCH_CMAKE_PREFIX_PATH "import torch; print(torch.utils.cmake_prefix_path)")
+        list(APPEND CMAKE_PREFIX_PATH "${TORCH_CMAKE_PREFIX_PATH}")
+        find_package(Torch REQUIRED CONFIG)
+    endif()
+    message(STATUS "Warp Torch SDK: ${Torch_DIR}")
     add_library(leapp_wrp_torch_custom_op SHARED torch/torch_warp_runner_op.cc)
     target_link_libraries(leapp_wrp_torch_custom_op PRIVATE leapp_warp_runtime_core "${TORCH_LIBRARIES}")
     target_include_directories(leapp_wrp_torch_custom_op PRIVATE "${TORCH_INCLUDE_DIRS}")
     if(NOT WIN32)
-        set_target_properties(leapp_wrp_torch_custom_op PROPERTIES
-            BUILD_RPATH "${WARP_BIN_DIR};${CUDAToolkit_LIBRARY_DIR}"
-            INSTALL_RPATH "${WARP_BIN_DIR};${CUDAToolkit_LIBRARY_DIR}"
-        )
+        set_property(TARGET leapp_wrp_torch_custom_op APPEND PROPERTY
+            BUILD_RPATH "${WARP_BIN_DIR};${CUDAToolkit_LIBRARY_DIR}")
     endif()
 endif()
+
+foreach(ADAPTER IN ITEMS leapp_wrp_onnx_custom_op leapp_wrp_torch_custom_op)
+    if(TARGET ${ADAPTER})
+        install(TARGETS ${ADAPTER}
+            LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+            RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR})
+    endif()
+endforeach()

leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/README.md

@@ -15,8 +15,10 @@ leapp-build-warp-runtime
 leapp-build-warp-runtime --status
 ```
 
-The build command validates existing dependencies without installing them,
-then builds both adapters under the standard user cache:
+The build command runs CMake to find and validate existing SDKs, then builds both
+adapters under the standard user cache. Before replacing an existing build, it
+validates the SDKs in a temporary CMake directory; failed validation leaves the
+cached build untouched. It never installs dependencies:
 
 - Linux: `build/libleapp_wrp_onnx_custom_op.so` and
   `build/libleapp_wrp_torch_custom_op.so`
@@ -25,9 +27,55 @@ then builds both adapters under the standard user cache:
 
 LEAPP discovers the cached libraries automatically. Use
 `LEAPP_WARP_ONNX_CUSTOM_OP_LIBRARY` and
-`LEAPP_WARP_PT2_CUSTOM_OP_LIBRARY` only to override the cache paths. Build with
-the Python environment that will load the libraries because the adapters link
-against that environment's Warp and PyTorch installations.
+`LEAPP_WARP_PT2_CUSTOM_OP_LIBRARY` to select other build or installation paths.
+By default, SDK discovery uses installed CUDA, CMake's Torch search paths, and
+Python packages from the interpreter running the helper.
+
+Pass CMake configuration after `--` to select explicit SDKs:
+
+```bash
+leapp-build-warp-runtime --build-dir /path/to/build -- \
+  -DWARP_ROOT=/path/to/warp \
+  -DTorch_DIR=/path/to/libtorch/share/cmake/Torch \
+  -DCUDAToolkit_ROOT=/usr/local/cuda \
+  -DORT_INCLUDE_DIR=/path/to/onnxruntime/include
+```
+
+`WARP_ROOT` contains `native/` and `bin/`; `ORT_INCLUDE_DIR` contains the ONNX
+Runtime C API headers. Explicit SDK locations take precedence over Python
+package discovery. Use Warp headers, `warp.so`, and exported graphs from the
+same Warp release, and a Torch SDK matching the process that loads the adapter.
+
+When no ONNX header directory is supplied, CMake downloads the release selected
+by `-DLEAPP_WARP_ONNXRUNTIME_VERSION=VERSION`, defaulting to Python's ONNX Runtime
+version. Downloads are cached by version. The adapter does not link another
+ONNX Runtime library; its host supplies the API.
+
+Install the adapters with standard CMake commands:
+
+```bash
+cmake --install /path/to/build --prefix /path/to/install
+```
+
+Libraries install under `CMAKE_INSTALL_LIBDIR` (`GNUInstallDirs`). Warp, Torch,
+and CUDA remain in their SDK locations. Build outputs use the discovered library
+paths; installed Linux adapters default to `$ORIGIN`. Configure
+`CMAKE_INSTALL_RPATH` or set `LD_LIBRARY_PATH` to the selected SDK library
+directories when loading installed adapters. Only the LEAPP adapters are installed.
+
+ONNX registers CPU and CUDA kernels. CUDA replay keeps data tensors on the GPU
+and uses ONNX Runtime's compute stream; only the embedded archive is copied to
+host once during loading. CPU replay stages data through CUDA device 0.
+Replay currently supports CUDA device 0 only and restores the caller's CUDA
+context. Calls sharing a graph are ordered across streams; warm each instance
+serially before concurrent use. Torch makes strided inputs contiguous before
+replay.
+
+Torch caches runners without retaining model constants. Expired entries are
+released on subsequent calls; native model owners can call
+`LEAPP_WARP_ReleaseUnusedRunners()` after releasing constants to clean up
+immediately (zero indicates success). Destruction waits for GPU work before
+freeing graph resources. Bundle tensors must remain immutable during inference.
 
 Windows builds also require a Warp import library. The `warp-lang` wheel does
 not currently include `warp.lib`; set `LEAPP_WARP_IMPORT_LIBRARY` to an