diff --git a/CMakeLists.txt b/CMakeLists.txt index a1e5a1e7..805feceb 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -16,6 +16,15 @@ cmake_minimum_required(VERSION 3.18 FATAL_ERROR) set(ACCELERATOR "cuda" CACHE STRING "Accelerator platform: cuda, metax, ascend, tsingmicro, dcu, gcu, or musa") +# Directory inside the wheel holding a bundled forked libtorch, when the backend +# ships one (see scripts/bundle_*_libtorch.sh). "lib" means "no separate bundle +# dir": the CUDA backend drops its extra .so straight into torch_fl/lib/. +# Set per backend below; consumed by cmake/FlagosRpath.cmake. +set(FLAGOS_BUNDLE_LIBDIR "lib") +# Vendor driver/runtime dirs to embed in RUNPATH. Never bundled into the wheel +# (a machine with the card has the driver), mirroring MetaX's /opt/maca. +set(FLAGOS_VENDOR_RPATH_DIRS "") + # Kernel build options (can be overridden via -D flags from setup.py) option(CUDA_KERNEL "Build CUDA kernel implementations" ON) option(FLAGGEMS_KERNEL "Build FlagGems C++ kernel wrappers (requires liboperators.so)" ON) @@ -54,6 +63,9 @@ elseif(ACCELERATOR STREQUAL "metax") elseif(NOT FLAGGEMS_KERNEL) set(FLAGGEMS_KERNEL OFF CACHE BOOL "Build FlagGems kernel implementations" FORCE) endif() + # Self-contained wheel: forked libtorch bundled by + # scripts/bundle_maca_libtorch.sh (the maca runtime stays on the target). + set(FLAGOS_BUNDLE_LIBDIR "lib_maca") project(TORCH_FLAGOS CXX C) elseif(ACCELERATOR STREQUAL "tsingmicro") project(TORCH_FLAGOS CXX C) @@ -67,6 +79,10 @@ elseif(ACCELERATOR STREQUAL "dcu") # No nvcc and no CUDA language: nothing outside backends/metax/ is a .cu. set(CUDA_KERNEL OFF CACHE BOOL "Build CUDA kernel implementations" FORCE) set(FLAGGEMS_KERNEL OFF CACHE BOOL "Build FlagGems kernel implementations" FORCE) + # Self-contained wheel: DTK's forked libtorch (libtorch_cpu.so carries hip + # symbols and needs libgalaxyhip.so.5) is bundled here by + # scripts/bundle_dcu_libtorch.sh. The DTK driver itself stays on the target. + set(FLAGOS_BUNDLE_LIBDIR "lib_dcu") project(TORCH_FLAGOS CXX C) elseif(ACCELERATOR STREQUAL "gcu") project(TORCH_FLAGOS CXX C) @@ -80,6 +96,15 @@ elseif(ACCELERATOR STREQUAL "musa") set(FLAGGEMS_KERNEL OFF CACHE BOOL "Build FlagGems kernel implementations" FORCE) project(TORCH_FLAGOS CXX C) else() + # PPU rides this branch: it has no accelerator value of its own because it is + # built against PPU_SDK/CUDA_SDK, so the CUDA boxing kernels apply unchanged + # (see the README's PPU section). It still needs its own bundle dir, since that + # libtorch is a local PPU build and not the upstream one. Detected by the SDK + # env vars, matching setup.py::_vendor_supplies_triton(). + if(DEFINED ENV{PPU_SDK} OR DEFINED ENV{PPU_HOME}) + set(FLAGOS_BUNDLE_LIBDIR "lib_ppu") + message(STATUS "PPU SDK detected: bundling libtorch into torch_fl/lib_ppu") + endif() # project(... CUDA) runs nvcc compiler-id before find_package(CUDAToolkit); conda # layouts need CUDAToolkit_ROOT (headers under targets/x86_64-linux/include). if(NOT CUDAToolkit_ROOT) @@ -222,7 +247,23 @@ elseif(ACCELERATOR STREQUAL "dcu") set(CUDA_RUNTIME_LIB cuda_runtime_compat) # Embed the compat + DTK library paths so LD_LIBRARY_PATH is not needed. - list(APPEND CMAKE_INSTALL_RPATH "${DCU_CUDA_ROOT}/lib64" "${DTK_ROOT}/lib") + # Measured layout of the deps the forked libtorch actually pulls. They are all + # on the DTK container's LD_LIBRARY_PATH, which a clean env does NOT inherit -- + # so the RUNPATH has to name them or a self-contained wheel gets "not found": + # cuda/cuda-*/lib64 libcudart.so.12 (shim over libgalaxyhip) + # lib libgalaxyhip.so.5 libhipnn librocfft.so.0 + # librocrand.so.1 librocsparse.so.1 libunwind.so.8 + # libMIOpen-recommend.so + # hip/lib, lib64 hip runtime + # aillvm/lib libomp.so + # .hyhal/rocm_smi/lib librocm_smi64.so.2 + # Keep in sync with VENDOR_RPATH in scripts/bundle_dcu_libtorch.sh. + set(_dcu_rpath_dirs + "${DCU_CUDA_ROOT}/lib64" "${DTK_ROOT}/lib" "${DTK_ROOT}/hip/lib" + "${DTK_ROOT}/lib64" "${DTK_ROOT}/aillvm/lib" "${DTK_ROOT}/llvm/lib" + "${DTK_ROOT}/.hyhal/rocm_smi/lib" "/opt/hyhal/lib") + list(APPEND CMAKE_INSTALL_RPATH ${_dcu_rpath_dirs}) + set(FLAGOS_VENDOR_RPATH_DIRS ${_dcu_rpath_dirs}) # Pre-create torch::cudart so PyTorch's cuda.cmake is skipped entirely # (guard: if(TARGET torch::cudart) return()). @@ -358,6 +399,12 @@ elseif(ACCELERATOR STREQUAL "metax") endif() message(STATUS "MetaX SDK path: ${METAX_PATH}") + # The maca runtime (mcblas/mcdnn/...) comes from the target's /opt/maca driver + # install; it is never bundled. /opt/maca is kept as an explicit fallback in + # case the build machine's METAX_PATH is a non-default location. + set(FLAGOS_VENDOR_RPATH_DIRS + "${METAX_PATH}/lib" "${METAX_PATH}/lib64" "/opt/maca/lib" "/opt/maca/lib64") + add_library(cuda_runtime_compat INTERFACE IMPORTED) set_target_properties(cuda_runtime_compat PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${METAX_PATH}/tools/cu-bridge/include;${METAX_PATH}/include;${METAX_PATH}/include/mcr" @@ -513,11 +560,15 @@ else() endif() if(PYTORCH_INSTALL_DIR) + # Needed for the *build* link step (and for an in-place/editable install, where + # the built .so stay in the source tree). flagos_set_portable_rpath() strips it + # from the installed targets' RUNPATH so the wheel is relocatable. list(APPEND CMAKE_INSTALL_RPATH "${PYTORCH_INSTALL_DIR}/lib") endif() include_directories(${CMAKE_CURRENT_SOURCE_DIR}) include(${PROJECT_SOURCE_DIR}/cmake/TorchPythonTargets.cmake) +include(${PROJECT_SOURCE_DIR}/cmake/FlagosRpath.cmake) add_subdirectory(${PROJECT_SOURCE_DIR}/csrc/runtime/accelerator) add_subdirectory(${PROJECT_SOURCE_DIR}/csrc) diff --git a/cmake/FlagosRpath.cmake b/cmake/FlagosRpath.cmake new file mode 100644 index 00000000..d408e519 --- /dev/null +++ b/cmake/FlagosRpath.cmake @@ -0,0 +1,76 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Portable RUNPATH for the installed native libs. +# +# The problem: CMakeLists.txt appends "${PYTORCH_INSTALL_DIR}/lib" to +# CMAKE_INSTALL_RPATH so an in-place build finds the active torch wheel. That is +# the *build machine's* interpreter path (e.g. +# /nfs/.../envs/torch-fl-210/lib/python3.12/site-packages/torch/lib), which does +# not exist on a target machine -- and worse, on a machine that does have that +# path it silently wins over the interpreter actually running, so a py3.12 env +# can load another env's libtorch. +# +# Dropping it is safe: libtorch_fl.so is only ever loaded through +# `import torch_fl` -> torch_fl._C -> libtorch_bindings.so, which happens after +# torch_fl/__init__.py has already run `import torch`. libc10 / libtorch_cpu are +# mapped by then, so the loader satisfies those DT_NEEDED entries by soname from +# the already-loaded set. This is why the MetaX self-contained wheel works today. +# +# Anything else the backend put on CMAKE_INSTALL_RPATH (vendor driver dirs, the +# FlagGems liboperators dir) is preserved, so in-place builds keep working. +# +# flagos_set_portable_rpath( [EXTRA_DIRS ...]) +function(flagos_set_portable_rpath _target) + cmake_parse_arguments(_ARG "" "" "EXTRA_DIRS" ${ARGN}) + + if(NOT UNIX OR APPLE) + return() + endif() + if(NOT TARGET ${_target}) + return() + endif() + + # $ORIGIN is torch_fl/lib/ for every installed target. A sibling bundle dir + # (torch_fl/lib_maca, lib_dcu, lib_ppu) is reached via $ORIGIN/../. + set(_rpath "$ORIGIN" "$ORIGIN/lib") + if(FLAGOS_BUNDLE_LIBDIR AND NOT FLAGOS_BUNDLE_LIBDIR STREQUAL "lib") + list(APPEND _rpath "$ORIGIN/../${FLAGOS_BUNDLE_LIBDIR}") + endif() + + foreach(_dir IN LISTS _ARG_EXTRA_DIRS FLAGOS_VENDOR_RPATH_DIRS) + if(_dir) + list(APPEND _rpath "${_dir}") + endif() + endforeach() + + # Inherit whatever the backend branches appended, minus the build machine's + # torch/lib and minus the $ORIGIN entries already placed above. + foreach(_dir IN LISTS CMAKE_INSTALL_RPATH) + if(NOT _dir MATCHES "^\\$ORIGIN" + AND NOT (PYTORCH_INSTALL_DIR AND _dir STREQUAL "${PYTORCH_INSTALL_DIR}/lib")) + list(APPEND _rpath "${_dir}") + endif() + endforeach() + + list(REMOVE_DUPLICATES _rpath) + string(REPLACE ";" ":" _rpath_str "${_rpath}") + + set_target_properties(${_target} PROPERTIES + INSTALL_RPATH "${_rpath_str}" + # Stop CMake appending the imported torch target's build-machine link dir. + INSTALL_RPATH_USE_LINK_PATH OFF + BUILD_WITH_INSTALL_RPATH ON + ) +endfunction() diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index 520e2cd6..af5c4f01 100644 --- a/csrc/CMakeLists.txt +++ b/csrc/CMakeLists.txt @@ -325,29 +325,18 @@ if(ACCELERATOR STREQUAL "metax") message(WARNING "mcblas not found, MetaX mm will fallback to naive kernel") endif() - # Self-contained wheel: the installed libtorch_fl.so must NOT bake the build - # machine's absolute torch/lib path into its RUNPATH (that path does not exist - # on target machines). Resolve the forked libtorch bundled under - # torch_fl/lib_maca/ via $ORIGIN, and the maca runtime (mcblas/mcdnn/...) from - # the target's /opt/maca driver install. Disabling INSTALL_RPATH_USE_LINK_PATH - # stops CMake from appending the imported torch target's build-machine link - # dir. Mirrors scripts/bundle_maca_libtorch.sh's patchelf pass. - set(_metax_rpath - "$ORIGIN:$ORIGIN/../lib_maca:${METAX_PATH}/lib:${METAX_PATH}/lib64:/opt/maca/lib:/opt/maca/lib64") - # Setting INSTALL_RPATH explicitly discards the entries the top-level - # CMakeLists appended to CMAKE_INSTALL_RPATH, including FlagGems' - # liboperators.so dir. Re-add it so the C++ FlagGems path (FLAGGEMS_KERNEL, - # kFlagOs) resolves liboperators.so without LD_LIBRARY_PATH. - if(FLAGGEMS_KERNEL AND _flaggems_libdir) - set(_metax_rpath "${_metax_rpath}:${_flaggems_libdir}") - endif() - set_target_properties(${LIBRARY_NAME} PROPERTIES - INSTALL_RPATH "${_metax_rpath}" - INSTALL_RPATH_USE_LINK_PATH OFF - BUILD_WITH_INSTALL_RPATH ON - ) endif() +# Self-contained wheel: the installed libtorch_fl.so must NOT bake the build +# machine's absolute torch/lib path into its RUNPATH -- that path does not exist +# on a target machine, and on one where it does exist it silently overrides the +# interpreter actually running. Resolve a bundled forked libtorch via $ORIGIN +# instead. Everything else the backend put on CMAKE_INSTALL_RPATH is kept, +# including FlagGems' liboperators.so dir, so the C++ FlagGems path +# (FLAGGEMS_KERNEL, kFlagOs) still resolves without LD_LIBRARY_PATH. Mirrors the +# patchelf pass in scripts/bundle_*_libtorch.sh. +flagos_set_portable_rpath(${LIBRARY_NAME}) + install(TARGETS ${LIBRARY_NAME} ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} diff --git a/csrc/aten/backends/flagos/python_op_caller.cc b/csrc/aten/backends/flagos/python_op_caller.cc index 9c196342..b8eb7aea 100644 --- a/csrc/aten/backends/flagos/python_op_caller.cc +++ b/csrc/aten/backends/flagos/python_op_caller.cc @@ -610,6 +610,15 @@ at::Generator GetFlagosDefaultCudaGenerator(int64_t device_index) { } py::gil_scoped_acquire gil; py::module_ torch_cuda = py::module_::import("torch.cuda"); + // torch.cuda.default_generators is populated by CUDA lazy-init, not at import. + // A self-contained wheel front-ends a stock torch+cpu whose CUDA state nothing + // else touches, so the tuple is still empty here and indexing it would raise + // IndexError -- which surfaces as a failure of every device-side RNG op + // (randn/rand/normal_) while plain factories (empty/zeros/ones) work fine. + // Force lazy-init first; it is idempotent and cheap once initialized. + if (py::len(torch_cuda.attr("default_generators")) == 0) { + torch_cuda.attr("init")(); + } py::object gens = torch_cuda.attr("default_generators"); py::object py_gen = gens[py::cast(device_index)]; // torch.Generator -> at::Generator via THPGenerator unpack. diff --git a/csrc/aten/copy_ops.cc b/csrc/aten/copy_ops.cc index eb56e719..1f8f7778 100644 --- a/csrc/aten/copy_ops.cc +++ b/csrc/aten/copy_ops.cc @@ -26,8 +26,16 @@ // On the CUDA-family backends (including MetaX boxing) the flagos device shares // the vendor's CUDA streams, so the current stream is readable from c10::cuda. +// +// USE_DCU is excluded for the same reason runtime/guard.h:22 excludes it: the +// DCU wheel is hipified, so resolves through DTK's +// CUDA-compat shim against USE_ROCM-hipified torch headers and expands to hip* +// symbols the shim never declares (`'hipStreamCaptureStatus' was not declared`), +// and c10::cuda::getCurrentCUDAStream would not link there anyway -- DTK exports +// c10::hip with zero c10::cuda symbols. DCU still shares the vendor's streams, +// so it needs *some* barrier; see SyncCurrentStreamBeforeBlockingCopy below. #if !defined(USE_ASCEND) && !defined(USE_TSINGMICRO) && !defined(USE_GCU) && \ - !defined(USE_MUSA) + !defined(USE_MUSA) && !defined(USE_DCU) #define FLAGOS_COPY_HAS_CUDA_STREAM 1 #include #endif @@ -53,6 +61,13 @@ inline void SyncCurrentStreamBeforeBlockingCopy() { if (stream.stream() != nullptr) { stream.synchronize(); } +#elif defined(USE_DCU) + // DCU shares the vendor's streams, so it needs this barrier just as much as + // the other CUDA-family backends -- but it cannot ask which stream is current + // (no c10::cuda symbols in a hipified wheel; see the guard above). Fall back to + // a device-wide sync: a superset of the per-stream wait, so still correct, + // just coarser. Only reached on the blocking-copy paths. + ::DeviceSynchronize(); #endif } diff --git a/csrc/runtime/accelerator/CMakeLists.txt b/csrc/runtime/accelerator/CMakeLists.txt index ab285eae..998bea55 100644 --- a/csrc/runtime/accelerator/CMakeLists.txt +++ b/csrc/runtime/accelerator/CMakeLists.txt @@ -123,6 +123,11 @@ else() target_link_libraries(${LIBRARY_NAME} PRIVATE CUDA::cudart) endif() +# libflagos.so needs no torch symbols at all (only the vendor runtime), so the +# inherited ${PYTORCH_INSTALL_DIR}/lib RUNPATH entry is pure build-machine +# leakage. Vendor driver dirs are preserved. +flagos_set_portable_rpath(${LIBRARY_NAME}) + install(TARGETS ${LIBRARY_NAME} ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} diff --git a/scripts/bundle_dcu_libtorch.sh b/scripts/bundle_dcu_libtorch.sh new file mode 100644 index 00000000..4e52055c --- /dev/null +++ b/scripts/bundle_dcu_libtorch.sh @@ -0,0 +1,170 @@ +#!/usr/bin/env bash +# Bundle the DTK-forked libtorch C++ .so into torch_fl/lib_dcu/ for a +# self-contained single wheel. +# +# Why bundle the core libs (not just libtorch_hip.so): +# Measured symbol attribution: DTK's libtorch_cpu.so resolves 2101 undefined +# symbols in libtorch_fl.so, while libtorch_hip.so resolves 0 -- the fork is on +# the core-lib side (DTK's libtorch_cpu.so exports 128 hip symbols and carries +# DT_NEEDED: libgalaxyhip.so.5). libtorch_hip.so still has to be present: the +# boxing kernels require vendor kernels to be registered under the CUDA dispatch +# key. +# +# Why bundle torch.libs/: those 17 files are auditwheel-mangled common libs +# (glog/gflags/MKL/OpenMPI/hwloc/libxml2...), 5 of which are direct DT_NEEDED of +# DTK's libtorch_cpu.so. Measured, they carry 0 hip/dtk/rocm references and +# introduce no SDK binding; but the official torch+cpu wheel's torch.libs +# contains none of them (names carry a hash, system libs do not match either), +# so the wheel cannot run without bundling them. They sit in the same dir as the +# core libs and are resolved via $ORIGIN. +# +# Note on $ORIGIN semantics: glibc expands $ORIGIN from the path the object was +# *opened by*, not from the resolved real path. So opening libc10.so through the +# torch/lib symlink gives $ORIGIN = torch/lib, and it cannot find +# libgflags-8aee0f6c.so.2.1.2 that sits alongside it in lib_dcu. Runtime +# preloading therefore walks the bundle's original paths +# (_vendor_libtorch._preload_global), not the symlinks -- the same rationale is +# documented there. +# +# Not bundled: DTK driver stack (libgalaxyhip.so.5 libMIOpen.so.1 +# librocblas.so.4 libhipblas.so.2 librccl.so.1, etc., 12 sonames) stays on the +# target machine under /opt/dtk. SDK version binding already exists and is +# stronger (libtorch_hip.so's DT_NEEDED hard-codes librocblas.so.4); bundling +# that extra 129 MB would not tighten the binding. +# +# Usage: +# FLAGOS_DCU_TORCH_LIB= bash scripts/bundle_dcu_libtorch.sh +# DTK_ROOT=/opt/dtk bash scripts/bundle_dcu_libtorch.sh +# +# Should run after `python setup.py bdist_wheel` (ACCELERATOR=dcu) and before +# packing the wheel. Idempotent. + +set -euo pipefail + +REPO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +# shellcheck source=scripts/lib/bundle_common.sh +source "${REPO_DIR}/scripts/lib/bundle_common.sh" + +LIB_DCU="${REPO_DIR}/torch_fl/lib_dcu" +TORCH_FL_LIB="${REPO_DIR}/torch_fl/lib" +DTK_ROOT="${DTK_ROOT:-${ROCM_PATH:-/opt/dtk}}" + +SRC="${FLAGOS_DCU_TORCH_LIB:-}" +if [ -z "${SRC}" ]; then + SRC="$(bundle_find_vendor_torch_lib libtorch_hip.so dtk hip das || true)" +fi + +if [ -z "${SRC}" ] || [ ! -d "${SRC}" ]; then + echo "error: DTK torch/lib not found. Set FLAGOS_DCU_TORCH_LIB=" >&2 + exit 1 +fi +if [ ! -f "${SRC}/libtorch_hip.so" ]; then + echo "error: ${SRC} does not contain libtorch_hip.so, not a DTK torch/lib" >&2 + exit 1 +fi + +bundle_require_patchelf + +# Self-consistent set aligned with _dcu_libtorch_link._CORE_SO + _HIP_SO. +# libshm.so: direct DT_NEEDED of libtorch_python.so (torch.multiprocessing's +# shared-memory manager); DTK's torch/lib has the real file, but it is not on +# the symlink manifest, so it must be bundled for $ORIGIN resolution. +# libcaffe2_nvrtc.so: torch.cuda.init() dlopens it. DTK torch ships this file, +# stock +cpu wheel does not, so without bundling it the on-demand init inside +# GetFlagosDefaultCudaGenerator dies on "Error in dlopen: libcaffe2_nvrtc.so" +# (the RNG fix exposed it). +CORE_SO=(libc10.so libtorch_cpu.so libtorch.so libtorch_global_deps.so libtorch_python.so libcaffe2_nvrtc.so) +HIP_SO=(libc10_hip.so libtorch_hip.so libmagma.so libshm.so) + +# DTK driver stack measured layout (container LD_LIBRARY_PATH matches find results): +# lib/ libhipnn librocfft.so.0 librocrand.so.1 librocsparse.so.1 +# libMIOpen-recommend.so libunwind.so.8 libgalaxyhip.so.5 +# hip/lib/ hip runtime +# aillvm/lib/ libomp.so +# .hyhal/rocm_smi/lib/ librocm_smi64.so.2 +# All stay on the target machine (a box with DCU cards has /opt/dtk), but RPATH +# must cover them, else not-found when LD_LIBRARY_PATH is unset. +VENDOR_RPATH="${DTK_ROOT}/lib:${DTK_ROOT}/hip/lib:${DTK_ROOT}/lib64" +VENDOR_RPATH="${VENDOR_RPATH}:${DTK_ROOT}/aillvm/lib:${DTK_ROOT}/.hyhal/rocm_smi/lib" +VENDOR_RPATH="${VENDOR_RPATH}:${DTK_ROOT}/llvm/lib:/opt/hyhal/lib" + +echo "Source DTK torch/lib : ${SRC}" +echo "Target lib_dcu : ${LIB_DCU}" +echo "DTK driver path : ${DTK_ROOT}" + +# Libs inside the bundle must be openable from two paths: +# 1. Directly from lib_dcu/ (runtime preload walks this one, see +# _vendor_libtorch._preload_global) +# 2. Through the torch/lib/ symlinks (when `import torch` loads +# libtorch_global_deps.so itself) +# glibc expands $ORIGIN from the path the object was opened by, so path #2 gives +# $ORIGIN = torch/lib and cannot find libmpi-3fcb240d.so.40.40.3 and other +# auditwheel-mangled libs sitting in lib_dcu. Both dirs are siblings under +# site-packages (torch/lib -> ../../torch_fl/lib_dcu), so adding one more +# relative path covers both cases. Measured: without this, `import torch` before +# `import torch_fl` dies on libmpi not found. +BUNDLE_ORIGIN="\$ORIGIN:\$ORIGIN/../../torch_fl/lib_dcu" + +bundle_copy_so "${SRC}" "${LIB_DCU}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 1 "${CORE_SO[@]}" +bundle_copy_so "${SRC}" "${LIB_DCU}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 0 "${HIP_SO[@]}" + +# torch.libs/: auditwheel dir sibling to torch/, filenames carry hash suffixes, +# can only glob. +TORCH_LIBS="$(cd "${SRC}/../.." && pwd)/torch.libs" +if [ -d "${TORCH_LIBS}" ]; then + echo "Source torch.libs : ${TORCH_LIBS}" + _names=() + while IFS= read -r f; do + _names+=("$(basename "${f}")") + done < <(find "${TORCH_LIBS}" -maxdepth 1 -type f -name '*.so*' | sort) + if [ ${#_names[@]} -gt 0 ]; then + bundle_copy_so "${TORCH_LIBS}" "${LIB_DCU}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 0 "${_names[@]}" + else + echo "warning: ${TORCH_LIBS} has no .so, skipping" >&2 + fi +else + echo "warning: ${TORCH_LIBS} not found; the target will lack hash-suffixed" >&2 + echo " common libs (libglog-*.so.0 etc.) from libtorch_cpu.so's DT_NEEDED." >&2 +fi + +# libtorch_fl.so / libflagos.so also need DTK's CUDA compatibility layer +# libcudart.so.12 (cuda_runtime_compat, see CMakeLists.txt's DCU_CUDA_ROOT). +# cmake already wrote it into RUNPATH; cannot drop it when rewriting here -- +# else libcudart.so.12 becomes not-found in a clean environment. +_DCU_CUDA_LIB64="" +for _c in "${DTK_ROOT}"/cuda/cuda-*/lib64; do + if [ -f "${_c}/libcudart.so.12" ]; then + _DCU_CUDA_LIB64="${_c}" + break + fi +done +if [ -z "${_DCU_CUDA_LIB64}" ]; then + echo "warning: ${DTK_ROOT}/cuda/cuda-*/lib64 has no libcudart.so.12" >&2 +fi +PLUGIN_RPATH="\$ORIGIN:\$ORIGIN/../lib_dcu" +[ -n "${_DCU_CUDA_LIB64}" ] && PLUGIN_RPATH="${PLUGIN_RPATH}:${_DCU_CUDA_LIB64}" +bundle_rewrite_plugin_rpath "${TORCH_FL_LIB}" "${PLUGIN_RPATH}:${VENDOR_RPATH}" + +# torch/version.py is pure Python generated at build time, so swapping .so +# cannot change it. In the self-contained DCU wheel the base is stock torch+cpu, +# whose torch.version.hip reports None, and triton's hcu backend gates +# is_active() on that value (backends/hcu/driver.py is_active(): +# torch.cuda.is_available() and torch.version.hip is not None). None means never +# activate, and any flag_gems op dies in triton's driver factory: "0 active +# drivers ([])". Carry the vendor torch's own version.py; at import time +# _restore_dcu_hip_version() reads back the hip/rocm strings. +_VENDOR_VERSION_PY="$(cd "${SRC}/.." && pwd)/version.py" +if [ -f "${_VENDOR_VERSION_PY}" ]; then + cp -fL "${_VENDOR_VERSION_PY}" "${LIB_DCU}/vendor_version.py" + echo "Copied vendor version.py -> lib_dcu/vendor_version.py" + grep -E "^\s*(hip|rocm)\s*(:|=)" "${LIB_DCU}/vendor_version.py" || true +else + echo "warning: ${_VENDOR_VERSION_PY} not found, triton hcu backend may not activate" >&2 +fi + +bundle_summary "${LIB_DCU}" +bundle_check_needed "${LIB_DCU}" \ + "${DTK_ROOT}/lib" "${DTK_ROOT}/hip/lib" "${DTK_ROOT}/lib64" \ + "${DTK_ROOT}/aillvm/lib" "${DTK_ROOT}/llvm/lib" \ + "${DTK_ROOT}/.hyhal/rocm_smi/lib" "${_DCU_CUDA_LIB64:-/nonexistent}" \ + "/opt/hyhal/lib" "/usr/lib64" "/usr/lib" "/usr/lib/x86_64-linux-gnu" diff --git a/scripts/bundle_maca_libtorch.sh b/scripts/bundle_maca_libtorch.sh index 9d1e8b5b..598a5ac7 100644 --- a/scripts/bundle_maca_libtorch.sh +++ b/scripts/bundle_maca_libtorch.sh @@ -1,84 +1,77 @@ #!/usr/bin/env bash -# 把沐曦 fork 的 libtorch C++ .so 打进 torch_fl/lib_maca/，做成自包含单 wheel。 +# Bundle the MetaX-forked libtorch C++ .so into torch_fl/lib_maca/ for a +# self-contained single wheel. # -# 官方 torch+cpu wheel 缺沐曦 fork 的 libtorch（带 at::maca::* / wcuda* 符号）， -# boxing 产物 libtorch_fl.so 又必须链接这些符号。本脚本从沐曦 torch wheel 拷贝那批 -# libtorch .so 到 torch_fl/lib_maca/，并用 patchelf 重写 RPATH： -# - lib_maca 内 libtorch .so -> $ORIGIN:/opt/maca/lib:/opt/maca/lib64 -# （运行期从目标机 /opt/maca 找 mcblas/mcdnn 等 maca runtime；不打包 runtime） -# - torch_fl/lib/libtorch_fl.so -> $ORIGIN:$ORIGIN/../lib_maca -# （从包内 lib_maca 找 fork libtorch，去掉构建机写死的绝对路径） +# The official torch+cpu wheel lacks the MetaX-forked libtorch (which exports +# at::maca::* / wcuda* symbols), and the boxing artifact libtorch_fl.so must +# link against those symbols. This script copies that batch of libtorch .so +# from the MetaX torch wheel into torch_fl/lib_maca/, and rewrites RPATH with +# patchelf: +# - libtorch .so inside lib_maca -> $ORIGIN:/opt/maca/lib:/opt/maca/lib64 +# (at runtime find mcblas/mcdnn and other maca runtime from /opt/maca on +# the target machine; the runtime is not bundled) +# - torch_fl/lib/libtorch_fl.so -> $ORIGIN:$ORIGIN/../lib_maca +# (find the forked libtorch from lib_maca inside the package, removing the +# absolute path hard-coded on the build machine) # -# maca runtime（libmcblas 等，~4.9G）不打包：装了沐曦卡的机器必有 /opt/maca 驱动。 +# maca runtime (libmcblas etc., ~4.9G) is not bundled: machines with MetaX +# cards installed must have the /opt/maca driver. # -# 用法: +# Usage: # FLAGOS_MACA_TORCH_LIB= bash scripts/bundle_maca_libtorch.sh -# MACA_PATH=/opt/maca bash scripts/bundle_maca_libtorch.sh # 覆盖 maca 路径 +# MACA_PATH=/opt/maca bash scripts/bundle_maca_libtorch.sh # override maca path # -# 应在 `python setup.py bdist_wheel`（ACCELERATOR=metax）之后、打 wheel 之前跑， -# 或跑完再重打 wheel。幂等。 +# Should run after `python setup.py bdist_wheel` (ACCELERATOR=metax) and before +# packing the wheel, or re-pack the wheel after running. Idempotent. set -euo pipefail REPO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +# shellcheck source=scripts/lib/bundle_common.sh +source "${REPO_DIR}/scripts/lib/bundle_common.sh" + LIB_MACA="${REPO_DIR}/torch_fl/lib_maca" TORCH_FL_LIB="${REPO_DIR}/torch_fl/lib" MACA_PATH="${MACA_PATH:-${METAX_HOME:-${MACA_HOME:-/opt/maca}}}" -# 沐曦 torch/lib 来源：显式 env，或从 conda 里找 +metax torch。 +# MetaX torch/lib source: explicit env, or find +metax torch from conda. SRC="${FLAGOS_MACA_TORCH_LIB:-}" if [ -z "${SRC}" ]; then - SRC=$(python - <<'PY' 2>/dev/null || true -import importlib.util, os -spec = importlib.util.find_spec("torch") -if spec and spec.submodule_search_locations: - lib = os.path.join(spec.submodule_search_locations[0], "lib") - ver = os.path.join(spec.submodule_search_locations[0], "version.py") - txt = open(ver).read() if os.path.isfile(ver) else "" - if ("metax" in txt or "maca" in txt) and os.path.exists(os.path.join(lib, "libtorch_cuda.so")): - print(lib) -PY -) + SRC="$(bundle_find_vendor_torch_lib libtorch_cuda.so metax maca || true)" fi if [ -z "${SRC}" ] || [ ! -d "${SRC}" ]; then - echo "error: 找不到沐曦 torch/lib。设 FLAGOS_MACA_TORCH_LIB=" >&2 + echo "error: MetaX torch/lib not found. Set FLAGOS_MACA_TORCH_LIB=" >&2 exit 1 fi if [ ! -f "${SRC}/libtorch_cuda.so" ]; then - echo "error: ${SRC} 里没有 libtorch_cuda.so，不是沐曦 torch/lib" >&2 + echo "error: ${SRC} does not contain libtorch_cuda.so, not a MetaX torch/lib" >&2 exit 1 fi -command -v patchelf >/dev/null 2>&1 || { echo "error: 需要 patchelf（pip install patchelf）" >&2; exit 1; } +bundle_require_patchelf -# 与 _metax_libtorch_link._CORE_SO + _CUDA_SO 对齐的自洽集合。 +# Self-consistent set aligned with _metax_libtorch_link._CORE_SO + _CUDA_SO. CORE_SO=(libc10.so libtorch_cpu.so libtorch.so libtorch_global_deps.so libtorch_python.so) -CUDA_SO=(libc10_cuda.so libtorch_cuda.so libtorch_cuda_linalg.so) +CUDA_SO=(libc10_cuda.so libtorch_cuda.so libtorch_cuda_linalg.so libshm.so) + +VENDOR_RPATH="${MACA_PATH}/lib:${MACA_PATH}/lib64" +# Same as DCU: libs inside the bundle must be openable from both lib_maca/ +# and torch/lib/ symlink paths. +BUNDLE_ORIGIN="\$ORIGIN:\$ORIGIN/../../torch_fl/lib_maca" -echo "源沐曦 torch/lib : ${SRC}" -echo "目标 lib_maca : ${LIB_MACA}" -echo "maca runtime 路径: ${MACA_PATH}/lib" -mkdir -p "${LIB_MACA}" +echo "Source MetaX torch/lib : ${SRC}" +echo "Target lib_maca : ${LIB_MACA}" +echo "maca runtime path : ${MACA_PATH}/lib" -for so in "${CORE_SO[@]}" "${CUDA_SO[@]}"; do - src="${SRC}/${so}" - if [ ! -f "${src}" ]; then - echo " 跳过 ${so}（源不存在）" - continue - fi - # 解引用软链，拷实体文件。 - cp -fL "${src}" "${LIB_MACA}/${so}" - patchelf --set-rpath "\$ORIGIN:${MACA_PATH}/lib:${MACA_PATH}/lib64" "${LIB_MACA}/${so}" - echo " 打包 ${so} ($(du -h "${LIB_MACA}/${so}" | cut -f1))" -done +bundle_copy_so "${SRC}" "${LIB_MACA}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 0 \ + "${CORE_SO[@]}" "${CUDA_SO[@]}" -# torch_fl.so 去掉构建机写死的沐曦 torch/lib 绝对路径，改为从包内 lib_maca 找 fork libtorch。 -for so in libtorch_fl.so libtorch_bindings.so; do - target="${TORCH_FL_LIB}/${so}" - [ -f "${target}" ] || continue - patchelf --set-rpath "\$ORIGIN:\$ORIGIN/../lib_maca:${MACA_PATH}/lib:${MACA_PATH}/lib64" "${target}" - echo " 重写 RPATH ${so}" -done +# Strip the build machine's hard-coded MetaX torch/lib absolute path from +# torch_fl.so, rewrite to find the forked libtorch from lib_maca inside the package. +bundle_rewrite_plugin_rpath "${TORCH_FL_LIB}" \ + "\$ORIGIN:\$ORIGIN/../lib_maca:${VENDOR_RPATH}" -echo "完成。lib_maca 总大小: $(du -sh "${LIB_MACA}" | cut -f1)" +bundle_summary "${LIB_MACA}" +bundle_check_needed "${LIB_MACA}" \ + "${MACA_PATH}/lib" "${MACA_PATH}/lib64" "/usr/lib64" "/usr/lib" diff --git a/scripts/bundle_ppu_libtorch.sh b/scripts/bundle_ppu_libtorch.sh new file mode 100644 index 00000000..12b15dd5 --- /dev/null +++ b/scripts/bundle_ppu_libtorch.sh @@ -0,0 +1,93 @@ +#!/usr/bin/env bash +# Bundle the locally built PPU libtorch C++ .so into torch_fl/lib_ppu/ for a +# self-contained single wheel. +# +# PPU is compiled against PPU_SDK/CUDA_SDK, so ACCELERATOR=cuda and the CUDA +# boxing kernels work as-is. The difference from a real NVIDIA machine lies in +# libtorch: it is a local USE_CUDA=1 source build, not an upstream wheel. +# Measured undefined symbols in libtorch_fl.so show that its libtorch_cpu.so +# provides 2092 of them, libtorch_cuda.so provides 0, and libc10_cuda.so +# provides 10 — so the core libs must be replaced, while libtorch_cuda.so just +# needs to be present (the CUDA dispatch key must have registered vendor kernels). +# +# That local build also links against system MKL in /usr/local/lib +# (libmkl_core / libmkl_gnu_thread / libmkl_intel_lp64, ~171 MB); the official +# torch+cpu wheel does not ship these files, so they are bundled into lib_ppu/ +# together. +# +# Not bundled: PPU SDK runtime stays on the target machine at /usr/local/PPU_SDK. +# +# Usage: +# FLAGOS_PPU_TORCH_LIB= bash scripts/bundle_ppu_libtorch.sh +# PPU_SDK=/usr/local/PPU_SDK bash scripts/bundle_ppu_libtorch.sh +# FLAGOS_PPU_MKL_DIR=/usr/local/lib bash scripts/bundle_ppu_libtorch.sh +# +# Should run after `python setup.py bdist_wheel` and before packing the wheel. +# Idempotent. + +set -euo pipefail + +REPO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +# shellcheck source=scripts/lib/bundle_common.sh +source "${REPO_DIR}/scripts/lib/bundle_common.sh" + +LIB_PPU="${REPO_DIR}/torch_fl/lib_ppu" +TORCH_FL_LIB="${REPO_DIR}/torch_fl/lib" +PPU_SDK="${PPU_SDK:-${PPU_HOME:-/usr/local/PPU_SDK}}" +MKL_DIR="${FLAGOS_PPU_MKL_DIR:-/usr/local/lib}" + +SRC="${FLAGOS_PPU_TORCH_LIB:-}" +if [ -z "${SRC}" ]; then + # PPU torch is locally built; version.py may not contain the "ppu" marker, + # so first try finding it by marker, fall back to "any torch with + # libtorch_cuda.so from the current interpreter" if not found. + SRC="$(bundle_find_vendor_torch_lib libtorch_cuda.so ppu || true)" + if [ -z "${SRC}" ]; then + SRC="$(bundle_find_vendor_torch_lib libtorch_cuda.so || true)" + fi +fi + +if [ -z "${SRC}" ] || [ ! -d "${SRC}" ]; then + echo "error: PPU torch/lib not found. Set FLAGOS_PPU_TORCH_LIB=" >&2 + exit 1 +fi +if [ ! -f "${SRC}/libtorch_cuda.so" ]; then + echo "error: ${SRC} does not contain libtorch_cuda.so, not a CUDA-built torch/lib" >&2 + exit 1 +fi + +bundle_require_patchelf + +# Self-consistent set aligned with _ppu_libtorch_link._CORE_SO + _CUDA_SO. +CORE_SO=(libc10.so libtorch_cpu.so libtorch.so libtorch_global_deps.so libtorch_python.so) +CUDA_SO=(libc10_cuda.so libtorch_cuda.so libtorch_cuda_linalg.so libshm.so) +MKL_SO=(libmkl_core.so.1 libmkl_gnu_thread.so.1 libmkl_intel_lp64.so.1) + +VENDOR_RPATH="${PPU_SDK}/CUDA_SDK/lib64:${PPU_SDK}/lib:${PPU_SDK}/lib64" +# Same as DCU: libs inside the bundle must be openable from both lib_ppu/ +# and torch/lib/ symlink paths. +BUNDLE_ORIGIN="\$ORIGIN:\$ORIGIN/../../torch_fl/lib_ppu" + +echo "Source PPU torch/lib : ${SRC}" +echo "Target lib_ppu : ${LIB_PPU}" +echo "PPU SDK path : ${PPU_SDK}" + +bundle_copy_so "${SRC}" "${LIB_PPU}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 1 "${CORE_SO[@]}" +bundle_copy_so "${SRC}" "${LIB_PPU}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 0 "${CUDA_SO[@]}" + +# System MKL: a direct DT_NEEDED of libtorch_cpu.so, with no same-named file +# in the official wheel. +if [ -d "${MKL_DIR}" ]; then + echo "Source MKL : ${MKL_DIR}" + bundle_copy_so "${MKL_DIR}" "${LIB_PPU}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 0 "${MKL_SO[@]}" +else + echo "warning: ${MKL_DIR} does not exist, skipping MKL; libs will be missing if the target machine has no MKL" >&2 +fi + +bundle_rewrite_plugin_rpath "${TORCH_FL_LIB}" \ + "\$ORIGIN:\$ORIGIN/../lib_ppu:${VENDOR_RPATH}" + +bundle_summary "${LIB_PPU}" +bundle_check_needed "${LIB_PPU}" \ + "${PPU_SDK}/CUDA_SDK/lib64" "${PPU_SDK}/lib" "${PPU_SDK}/lib64" \ + "${MKL_DIR}" "/usr/lib64" "/usr/lib" "/usr/local/lib" diff --git a/scripts/lib/bundle_common.sh b/scripts/lib/bundle_common.sh new file mode 100644 index 00000000..8cdf023a --- /dev/null +++ b/scripts/lib/bundle_common.sh @@ -0,0 +1,147 @@ +#!/usr/bin/env bash +# common parts for self-contained wheel bundling, sourced by scripts/bundle_{maca,dcu,ppu}_libtorch.sh. +# +# background: metax / dcu / ppu all run on a vendor fork of libtorch -- measured +# symbol attribution shows libtorch_cpu.so provides 2000+ undefined symbols for +# libtorch_fl.so, while the vendor lib (libtorch_hip.so / libtorch_cuda.so) +# provides 0. so a self-contained wheel must bundle the core libs together, not +# just the vendor lib. unified approach: +# - .so inside the bundle dir -> RPATH = $ORIGIN: +# (driver runtime stays on the target machine, not bundled: a box with the +# card already has the driver) +# - torch_fl/lib/libtorch_fl.so -> RPATH = $ORIGIN:$ORIGIN/../: +# (removes the build machine's hard-coded absolute paths. cmake/FlagosRpath.cmake +# already sets this; the patchelf here is fallback for existing .so / non-cmake +# build artifacts, kept idempotent.) +# +# the runtime logic that symlinks the bundle's core libs into stock torch/lib is +# in torch_fl/accelerator/_vendor_libtorch.py; the .so list on both sides must align. +# +# provides: +# bundle_require_patchelf +# bundle_find_vendor_torch_lib +# bundle_copy_so +# bundle_rewrite_plugin_rpath +# bundle_summary + +set -euo pipefail + +# patchelf is missing by default on all four nodes. on bw1000/810e `pip install patchelf` +# works directly; mc550's default index lacks the package and needs an explicit source. +bundle_require_patchelf() { + if command -v patchelf >/dev/null 2>&1; then + return 0 + fi + cat >&2 <<'EOF' +error: patchelf not found, cannot rewrite RPATH. install with (any): + pip install patchelf + pip install -i https://pypi.tuna.tsinghua.edu.cn/simple patchelf # when default index lacks it + conda install -c conda-forge patchelf +EOF + return 1 +} + +# locate the vendor torch/lib: prefer the current interpreter's torch (verify +# version.py contains a vendor marker and probe_so exists), return non-zero on +# failure. caller should check their own env override variable first. +# usage: bundle_find_vendor_torch_lib libtorch_hip.so dtk hip +bundle_find_vendor_torch_lib() { + local probe="$1" + shift + local markers="$*" + local py + py="$(command -v python || command -v python3 || true)" + [ -n "${py}" ] || return 1 + PROBE_SO="${probe}" VENDOR_MARKERS="${markers}" "${py}" - <<'PY' 2>/dev/null || return 1 +import importlib.util, os + +probe = os.environ["PROBE_SO"] +markers = os.environ["VENDOR_MARKERS"].split() +spec = importlib.util.find_spec("torch") +if spec and spec.submodule_search_locations: + root = spec.submodule_search_locations[0] + lib = os.path.join(root, "lib") + ver = os.path.join(root, "version.py") + txt = open(ver).read() if os.path.isfile(ver) else "" + if (not markers or any(m in txt for m in markers)) and os.path.exists( + os.path.join(lib, probe) + ): + print(lib) +PY +} + +# copy .so and set RPATH. cp -fL dereferences symlinks, copies the real file. +# when required=1 a missing source fails; =0 skips it (the stock +cpu wheel does +# not ship vendor libs anyway). +bundle_copy_so() { + local src_dir="$1" dst_dir="$2" rpath="$3" required="$4" + shift 4 + local so src + mkdir -p "${dst_dir}" + for so in "$@"; do + src="${src_dir}/${so}" + if [ ! -f "${src}" ]; then + if [ "${required}" = "1" ]; then + echo "error: missing required ${src}" >&2 + return 1 + fi + echo "warning: ${src} does not exist, not bundled" >&2 + continue + fi + cp -fL "${src}" "${dst_dir}/${so}" + # non-ELF (rare .so that are actually linker scripts) will fail patchelf, not fatal. + if ! patchelf --set-rpath "${rpath}" "${dst_dir}/${so}" 2>/dev/null; then + echo " ${so}: patchelf skipped (non-ELF?)" + fi + echo " bundled ${so} ($(du -h "${dst_dir}/${so}" | cut -f1))" + done +} + +# rewrite RPATH for libtorch_fl.so / libtorch_bindings.so / libflagos.so. +bundle_rewrite_plugin_rpath() { + local lib_dir="$1" rpath="$2" so target + for so in libtorch_fl.so libtorch_bindings.so libflagos.so; do + target="${lib_dir}/${so}" + [ -f "${target}" ] || continue + patchelf --set-rpath "${rpath}" "${target}" + echo " rewrote RPATH ${so} -> ${rpath}" + done +} + +bundle_summary() { + local dst_dir="$1" + echo "done. $(basename "${dst_dir}") total size: $(du -sh "${dst_dir}" | cut -f1) ($(find "${dst_dir}" -type f | wc -l | tr -d ' ') files)" +} + +# baseline libs guaranteed present on the target machine, not reported during self-check. +_BUNDLE_BASELINE_RE='^(libc|libm|libdl|librt|libpthread|libstdc\+\+|libgcc_s|ld-linux.*|libutil|libresolv|libnsl|libcrypt|libatomic|libgomp|libz|libnuma)\.so' + +# self-check: for each .so in the bundle, report every DT_NEEDED that is neither in +# the bundle, nor in the given target machine directories, nor a baseline system lib +# -- those are candidates for "not found" on the target machine. +# usage: bundle_check_needed +bundle_check_needed() { + local dst_dir="$1" + shift + local so dep d found + echo "---- DT_NEEDED self-check (only reports items potentially missing on target) ----" + local missing=0 + for so in "${dst_dir}"/*.so*; do + [ -f "${so}" ] || continue + while read -r dep; do + [ -n "${dep}" ] || continue + [[ "${dep}" =~ ${_BUNDLE_BASELINE_RE} ]] && continue + [ -e "${dst_dir}/${dep}" ] && continue + found=0 + for d in "$@"; do + if [ -e "${d}/${dep}" ]; then found=1; break; fi + done + [ "${found}" = "1" ] && continue + echo " $(basename "${so}") -> ${dep}" + missing=1 + done < <(patchelf --print-needed "${so}" 2>/dev/null || true) + done + if [ "${missing}" = "0" ]; then + echo " none (bundle + driver directories already cover all non-baseline deps)" + fi +} diff --git a/setup.py b/setup.py index de46c287..d5db7956 100644 --- a/setup.py +++ b/setup.py @@ -34,6 +34,17 @@ # "dcu", "gcu", or "musa" ACCELERATOR = os.environ.get("ACCELERATOR", "cuda").lower() +# Directory inside the wheel holding a bundled forked libtorch, for the backends +# that ship one (see scripts/bundle_*_libtorch.sh). "lib" means "no separate +# bundle dir": the CUDA backend drops its extra .so straight into torch_fl/lib/. +# Must match FLAGOS_BUNDLE_LIBDIR in CMakeLists.txt -- _C.so's RUNPATH has to +# reach the bundle or its auditwheel-mangled deps (libglog-*.so.0) go missing. +_BUNDLE_LIBDIR = {"metax": "lib_maca", "dcu": "lib_dcu"}.get(ACCELERATOR, "lib") +if _BUNDLE_LIBDIR == "lib" and ( + os.environ.get("PPU_SDK") or os.environ.get("PPU_HOME") +): + _BUNDLE_LIBDIR = "lib_ppu" + BASE_DIR = os.path.dirname(os.path.realpath(__file__)) # Only run cmake build for actual build commands, not metadata collection @@ -363,17 +374,23 @@ def build_deps(): ] ) elif ACCELERATOR == "dcu": - # Pure boxing build. The DCU torch wheel is a hipified build whose HIP - # kernels are registered under the CUDA dispatch key, so the generated + # Boxing build. The DCU torch wheel is a hipified build whose HIP kernels + # are registered under the CUDA dispatch key, so the generated # PrivateUse1 -> CUDA boxing kernels reach them with no hand-written - # kernels of our own. FLAGGEMS_KERNEL needs liboperators.so (not built - # for DTK); FLAGGEMS_PYTHON needs a DTK triton, which is a separate - # install -- both stay off unless explicitly requested below. + # kernels of our own. FLAGGEMS_KERNEL needs liboperators.so, which is not + # built for DTK, and stays off. + # + # FLAGGEMS_PYTHON defaults ON, same as metax/cuda: DTK ships a working + # triton (hcu backend) that flag_gems runs on, so the wheel compiles the + # FlagGems Python-path kernels too and the choice becomes a runtime one + # (FLAGOS_USE_FLAGGEMS -> backends_dcu_flaggems.conf). python_op_caller + # links torch_python_library, already in the link set, so this adds + # nothing to the wheel size. Set FLAGGEMS_PYTHON=0 for a slim pure-boxing + # build; the generic pass-through below honors that. cmake_args.extend( [ "-DCUDA_KERNEL=OFF", "-DFLAGGEMS_KERNEL=OFF", - "-DFLAGGEMS_PYTHON=OFF", "-DMETAX_KERNEL=OFF", "-DASCEND_KERNEL=OFF", ] @@ -517,7 +534,8 @@ def _bundle_cuda_assets() -> None: externally-supplied libtorch_cuda.so (CPU-only pip torch does not ship it). Historically this was LD_PRELOAD-ed by scripts/with_cuda_libtorch.sh; for a single self-contained wheel we bundle the assets and preload them from - torch_fl/__init__.py before `import torch` (see docs §约束1). CUDA only. + torch_fl/__init__.py before `import torch` (see that doc, constraint 1). + CUDA only. Set FLAGOS_SKIP_CUDA_ASSETS=1 to skip (e.g. a slim build for a machine that supplies libtorch_cuda.so out-of-band). @@ -604,19 +622,33 @@ def run(self): os.remove(os.path.join(dirpath, filename)) +def _extension_rpath_args(): + """RUNPATH for torch_fl._C: torch_fl/lib plus the bundle dir when separate. + + _C.so links libtorch_bindings.so out of torch_fl/lib, which in turn pulls the + bundled vendor libtorch and its auditwheel-mangled deps out of the bundle dir. + Without the second entry a self-contained wheel fails at import with e.g. + "libglog.so.0: cannot open shared object file". + """ + args = make_relative_rpath_args("lib") + if _BUNDLE_LIBDIR != "lib": + args += make_relative_rpath_args(_BUNDLE_LIBDIR) + return args + + def _extension_compile_args(): if IS_WINDOWS: # /NODEFAULTLIB makes sure we only link to DLL runtime # and matches the flags set for protobuf and ONNX - extra_link_args: list[str] = ["/NODEFAULTLIB:LIBCMT.LIB"] + [ - *make_relative_rpath_args("lib") - ] + extra_link_args: list[str] = [ + "/NODEFAULTLIB:LIBCMT.LIB" + ] + _extension_rpath_args() # /MD links against DLL runtime # and matches the flags set for protobuf and ONNX # /EHsc is about standard C++ exception handling extra_compile_args = ["/MD", "/FS", "/EHsc"] else: - extra_link_args = [*make_relative_rpath_args("lib")] + extra_link_args = _extension_rpath_args() extra_compile_args = [ "-Wall", "-Wextra", @@ -649,10 +681,20 @@ def _get_setup_kwargs(): "lib/*.dylib*", "lib/*.dll", "lib/*.lib", - # MetaX self-contained wheel: forked libtorch C++ .so bundled here so - # the process loads the MetaX C++ runtime without a separate metax - # torch wheel (see scripts/bundle_maca_libtorch.sh). + # Self-contained wheels: the vendor's forked libtorch C++ .so bundled + # here so the process loads that C++ runtime without a separate + # vendor torch wheel (see scripts/bundle_*_libtorch.sh, and + # torch_fl/accelerator/_vendor_libtorch.py for the relink at import). + # The trailing * matters for lib_dcu: DTK's auditwheel-mangled + # torch.libs deps end in a version suffix (libglog-6ed04f2c.so.0.0.0). "lib_maca/*.so*", + "lib_dcu/*.so*", + # DTK torch's own version.py, carried so _restore_dcu_hip_version() + # can hand triton's hcu backend the hip/rocm strings the stock +cpu + # torch in front does not have. Needed explicitly: the globs above + # only match *.so*. + "lib_dcu/vendor_version.py", + "lib_ppu/*.so*", # All backend configs, not just the default: runtime op-routing # configs selected via FLAGOS_USE_FLAGGEMS (backends_flaggems.conf) # and boxing modes via FLAGOS_BACKEND_CONFIG (backends_cuda.conf / @@ -663,11 +705,20 @@ def _get_setup_kwargs(): } version = "0.1.0" - if ACCELERATOR == "metax": - # Local version segment tags the wheel as a MetaX build (self-contained - # forked libtorch). Overridable via FLAGOS_WHEEL_LOCAL for a concrete - # MACA/driver version, e.g. FLAGOS_WHEEL_LOCAL=metax3.8.1. - local = os.environ.get("FLAGOS_WHEEL_LOCAL", "metax") + # A local version segment tags which vendor a self-contained wheel bundles a + # forked libtorch for. That bundle is SDK-version-bound whether we say so or + # not -- DTK's libtorch_hip.so has librocblas.so.4 written into its + # DT_NEEDED -- so making the binding visible in the filename is strictly + # better than leaving two incompatible wheels both called 0.1.0. Override + # with FLAGOS_WHEEL_LOCAL to pin the exact SDK, e.g. + # FLAGOS_WHEEL_LOCAL=metax3.8.1 / FLAGOS_WHEEL_LOCAL=dtk2604. + _default_local = {"metax": "metax", "dcu": "dtk"}.get(ACCELERATOR) + if _default_local is None and ( + os.environ.get("PPU_SDK") or os.environ.get("PPU_HOME") + ): + _default_local = "ppu" + local = os.environ.get("FLAGOS_WHEEL_LOCAL", _default_local) + if local: version = f"{version}+{local}" return dict( diff --git a/torch_fl/__init__.py b/torch_fl/__init__.py index 169f3ba3..4863c580 100644 --- a/torch_fl/__init__.py +++ b/torch_fl/__init__.py @@ -13,6 +13,7 @@ # limitations under the License. import os +import re import sys @@ -165,27 +166,83 @@ def _select_backend_config() -> None: ensure_cudart_shim() -# When reusing PyTorch's CUDA boxing kernels on MetaX with a stock +cpu torch -# wheel, the active wheel's torch/lib must point at the MetaX C++ runtime .so. -# This MUST run before `import torch` (afterwards libc10 is already mapped and -# relinking is too late). Gated on FLAGOS_METAX_BOXING=1; idempotent; no-op when -# torch already IS the MetaX wheel. -if os.environ.get("FLAGOS_METAX_BOXING", "0") == "1": - from torch_fl.accelerator.metax._metax_libtorch_link import ( - ensure_maca_libtorch_links, - ) - ensure_maca_libtorch_links() +def _relink_vendor_libtorch() -> None: + """Point the active torch wheel's torch/lib at this wheel's bundled libtorch. + + MetaX, DCU and PPU all run on a *forked* libtorch whose core .so + (libc10/libtorch_cpu/libtorch_python/...) differ from the upstream ones a + stock ``torch==X.Y.Z+cpu`` wheel ships. A self-contained wheel bundles them + under torch_fl/lib_{maca,dcu,ppu}/ and symlinks them over the stock files; + see torch_fl.accelerator._vendor_libtorch for why a ctypes preload alone is + not enough there. + + This MUST run before `import torch` -- afterwards libc10 is already mapped + and relinking is too late. Every backend's entry point is idempotent and a + no-op when its bundle dir is absent (a plain in-place build, where torch + already IS the vendor wheel), so this is safe to call unconditionally. + + MetaX: FLAGOS_METAX_BOXING=1 triggers relink unconditionally (for in-place + MetaX builds that want to test boxing). When accel=="metax" and the bundle + dir exists, relink regardless of the env var (self-contained wheel). DCU + and PPU have no native-kernel mode, so bundle-dir presence alone decides. + The CUDA backend is not here: the official +cpu wheel's core .so ARE the + upstream ones, so only the extra CUDA libs are missing and + _preload_cuda_assets() below handles those with ctypes. + """ + accel = _build_accelerator() + + if os.environ.get("FLAGOS_METAX_BOXING", "0") == "1": + from torch_fl.accelerator.metax._metax_libtorch_link import ( + ensure_maca_libtorch_links, + ) + + ensure_maca_libtorch_links() + return + + if accel == "metax": + from torch_fl.accelerator._vendor_libtorch import bundled_lib_dir + + if bundled_lib_dir("lib_maca", "libtorch_cuda.so"): + from torch_fl.accelerator.metax._metax_libtorch_link import ( + ensure_maca_libtorch_links, + ) + + ensure_maca_libtorch_links() + return + + if accel == "dcu": + from torch_fl.accelerator.dcu._dcu_libtorch_link import ( + ensure_dcu_libtorch_links, + ) + + ensure_dcu_libtorch_links() + return + + # PPU builds as ACCELERATOR=cuda (it targets PPU_SDK/CUDA_SDK), so the only + # distinguishing signal at import time is its own bundle dir. + if accel in ("cuda", ""): + from torch_fl.accelerator._vendor_libtorch import bundled_lib_dir + + if bundled_lib_dir("lib_ppu", "libtorch_cuda.so"): + from torch_fl.accelerator.ppu._ppu_libtorch_link import ( + ensure_ppu_libtorch_links, + ) + + ensure_ppu_libtorch_links() + + +_relink_vendor_libtorch() def _preload_cuda_assets() -> None: """Load the bundled CUDA .so into this process BEFORE `import torch`. - Hard constraint (docs/cpu_torch_external_libtorch_cuda.md §约束1): PyTorch - caches its CUDAHooks on first `import torch`. If libtorch_cuda.so is loaded - afterwards, device init fails with "Cannot initialize CUDA without ATen_cuda - library" even though the kernels register. So we ctypes-dlopen it here, at - the very top of torch_fl, before torch is imported. + Hard constraint (docs/cpu_torch_external_libtorch_cuda.md, constraint 1): + PyTorch caches its CUDAHooks on first `import torch`. If libtorch_cuda.so is + loaded afterwards, device init fails with "Cannot initialize CUDA without + ATen_cuda library" even though the kernels register. So we ctypes-dlopen it + here, at the very top of torch_fl, before torch is imported. libtorch_cuda.so has unresolved deps on the NVIDIA runtime libs (libcudart, libcublas, libcudnn, libnvshmem_host, ...) shipped by the pip nvidia-*-cu12 @@ -322,7 +379,6 @@ def _check_privateuse1_unclaimed() -> None: import torch # noqa: E402 - if sys.platform == "win32": from ._utils import _load_dll_libraries @@ -371,7 +427,6 @@ def _check_privateuse1_unclaimed() -> None: from . import flagos # noqa: E402 - torch.utils.rename_privateuse1_backend("flagos") torch._register_device_module("flagos", flagos) torch.utils.generate_methods_for_privateuse1_backend(for_storage=True) @@ -458,6 +513,43 @@ def _patched(increment, generator=None): pass +def _restore_dcu_hip_version() -> None: + """Set torch.version.hip/rocm for a self-contained DCU wheel. + + See the DCU branch of _patch_flaggems_codegen_config() for why this matters: + the bundled libtorch is DTK's HIP build, but torch/version.py comes from the + stock torch+cpu wheel in front and reports hip=None, which switches triton's + hcu backend off. scripts/bundle_dcu_libtorch.sh copies the vendor torch's own + version.py next to the bundled .so as vendor_version.py; read the strings + back from there. No-op when a real vendor torch is in front. + """ + import torch + + if getattr(torch.version, "hip", None): + return # a real DTK torch is in front; leave its values alone. + + hip_ver = os.environ.get("FLAGOS_DCU_HIP_VERSION", "").strip() + rocm_ver = "" + if not hip_ver: + ver_py = os.path.join(os.path.dirname(__file__), "lib_dcu", "vendor_version.py") + try: + with open(ver_py, encoding="utf-8") as f: + for line in f: + m = re.match(r"\s*hip\s*(?::[^=]*)?=\s*'([^']+)'", line) + if m: + hip_ver = m.group(1) + continue + m = re.match(r"\s*rocm\s*(?::[^=]*)?=\s*'([^']+)'", line) + if m: + rocm_ver = m.group(1) + except OSError: + return # not a bundled build (source checkout); nothing to restore. + if hip_ver: + torch.version.hip = hip_ver + if rocm_ver: + torch.version.rocm = rocm_ver + + def _patch_flaggems_codegen_config(): """ Configure FlagGems' vendor + torch.cuda shim for the flagos device. @@ -522,6 +614,18 @@ def _patch_flaggems_codegen_config(): # fallback. setdefault so an explicit GEMS_VENDOR still wins. if _build_accelerator() == "dcu" and os.environ.get("GEMS_VENDOR") != "ascend": os.environ.setdefault("GEMS_VENDOR", "hygon") + # torch.version is pure Python (torch/version.py), generated when the + # wheel is built -- swapping the bundled DTK .so files cannot change it. + # A self-contained DCU wheel therefore front-ends a stock torch+cpu whose + # torch.version.hip is None, while the DTK torch it replaces reports + # e.g. "6.3.26113". triton's hcu backend gates on exactly that value + # (backends/hcu/driver.py is_active(): torch.cuda.is_available() and + # torch.version.hip is not None), so with None the driver never activates + # and any flag_gems op dies in triton's driver factory with + # "0 active drivers ([]). There should only be one." Restore the attribute + # from the bundled libtorch's own version so triton sees a HIP torch, + # matching what the vendor wheel reported. + _restore_dcu_hip_version() return # --- Enflame GCU branch --- @@ -640,8 +744,14 @@ def _patched_cuda_device_init(self, device): _patch_cuda_device_context() # Initialize CUDA runtime only when FlagGems Python path needs it (CUDA backend ops). +# The check must be against the *build* backend, not torch.cuda.is_available(): +# a DCU/PPU self-contained wheel relinks a hipified/cuda libtorch into a stock +# +cpu torch, which makes is_available() return True even though the CUDA runtime +# libs are absent, and torch.cuda.init() would fail with "libcaffe2_nvrtc.so: not +# found". Only actual CUDA-backend builds need this init. if ( os.environ.get("FLAGOS_DISABLE_FLAGGEMS_PY", "0") != "1" + and _build_accelerator() in ("cuda", "") and torch.cuda.is_available() ): torch.cuda.init() diff --git a/torch_fl/accelerator/_vendor_libtorch.py b/torch_fl/accelerator/_vendor_libtorch.py new file mode 100644 index 00000000..8151481f --- /dev/null +++ b/torch_fl/accelerator/_vendor_libtorch.py @@ -0,0 +1,302 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Point the active (stock) torch wheel's ``torch/lib`` at a bundled vendor libtorch. + +Why this exists +--------------- +Several backends run on a *forked* libtorch: MetaX (``at::maca::*``), Hygon DCU +(DTK's hipified build -- its ``libtorch_cpu.so`` carries hip symbols and needs +``libgalaxyhip.so.5``), and PPU (a local ``USE_CUDA=1`` build). Measured symbol +attribution says the fork lives in the *core* libs: for both DCU and PPU, +``libtorch_cpu.so`` resolves ~2100 of ``libtorch_fl.so``'s undefined symbols while +the vendor lib (``libtorch_hip.so`` / ``libtorch_cuda.so``) resolves 0. So a +self-contained wheel must ship the core libs, not just the vendor one. + +Pure ``ctypes`` preloading does NOT work for core libs. The stock wheel's +``_C.so`` / ``libtorch_python.so`` carry an ``$ORIGIN`` RUNPATH that pulls the +upstream ``libc10.so`` back in *by full path*, so the process ends up with two +libc10 and dies in duplicate static init (``Key already registered ... +caffe2_report_cpu_memory_usage``). The robust fix is to make the physical files +that RUNPATH resolves to *be* the vendor ones: replace the stock wheel's +``torch/lib/`` with symlinks into the bundle dir. Originals move to +``torch/lib/_orig_backup/``, so the operation is fully reversible. + +The CUDA backend is the one exception and does not use this module: the official +``+cpu`` wheel's core libs *are* the upstream ones, so only the extra CUDA libs +are missing and a ctypes preload (``torch_fl.__init__._preload_cuda_assets``) +suffices. + +Callers must invoke this from ``torch_fl/__init__.py`` BEFORE ``import torch`` +(afterwards libc10 is already mapped and relinking is too late). Every entry +point here is idempotent. + +Note on ``$ORIGIN`` and symlinks: glibc expands ``$ORIGIN`` from the path the +object was *loaded by*, NOT from its resolved target. Opening +``torch/lib/libtorch_cpu.so`` (a symlink) therefore gives ``$ORIGIN`` = +``torch/lib``, where a bundle-internal dependency does not exist -- measured on +DCU, whose ``libc10.so`` needs the auditwheel-mangled ``libgflags-8aee0f6c.so`` +that ships in the bundle dir: + + ctypes.CDLL(".../torch/lib/libc10.so") -> libgflags-...so: not found + ctypes.CDLL(".../torch_fl/lib_dcu/libc10.so") -> OK + +So ``_preload_global`` dlopens the *bundle* paths, not the symlinks. That also +covers the symlinks: glibc keys loaded objects by (device, inode), and a symlink +shares both with its target, so a later lookup that resolves through +``torch/lib`` finds the object already mapped instead of re-opening it. +""" + +import ctypes +import importlib.util +import os +import sys + +# One flag per bundle dir: a process only ever relinks for its own backend, but +# keying by name keeps the module reentrant and makes the no-op cheap. +_done = set() +# dlopen handles kept alive for the process lifetime (see _preload_global). +_runtime_handles = [] + + +def active_torch_lib(): + """``torch/lib`` of the importable torch, WITHOUT importing torch.""" + spec = importlib.util.find_spec("torch") + if spec is None or not spec.submodule_search_locations: + return None + lib = os.path.join(spec.submodule_search_locations[0], "lib") + return lib if os.path.isdir(lib) else None + + +def bundled_lib_dir(bundle_dirname, probe_so): + """The bundle dir inside this wheel, if the bundling step actually ran. + + ``scripts/bundle__libtorch.sh`` copies the vendor libtorch .so into + ``torch_fl//``. When present this is the preferred source: + the target machine then needs only the official ``torch+cpu`` wheel plus the + vendor driver runtime, no vendor torch wheel at all. Absent (a plain + in-place/dev build) every entry point here becomes a no-op. + """ + # this file: torch_fl/accelerator/_vendor_libtorch.py -> torch_fl/ + pkg_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + libdir = os.path.join(pkg_root, bundle_dirname) + if os.path.isdir(libdir) and os.path.exists(os.path.join(libdir, probe_so)): + return libdir + return None + + +def _scan_sibling_envs(probe_so, vendor_markers): + """Fallback for multi-env dev setups: a sibling conda env's vendor torch. + + Matches on ``torch/version.py`` containing one of ``vendor_markers`` (e.g. + "metax"/"maca", "dtk"/"hip", "ppu") so we never pick up a stock wheel. + """ + if not vendor_markers: + return None + prefix = os.environ.get("CONDA_PREFIX") or os.path.dirname( + os.path.dirname(os.__file__) + ) + envs_root = os.path.dirname(prefix) # .../envs + if not os.path.isdir(envs_root): + return None + py = "python{}.{}".format(*sys.version_info[:2]) + try: + names = sorted(os.listdir(envs_root)) + except OSError: + return None + for name in names: + cand = os.path.join(envs_root, name, "lib", py, "site-packages", "torch") + libdir = os.path.join(cand, "lib") + ver_file = os.path.join(cand, "version.py") + if not os.path.isfile(ver_file) or not os.path.isdir(libdir): + continue + try: + with open(ver_file) as f: + txt = f.read() + except OSError: + continue + if any(m in txt for m in vendor_markers) and os.path.exists( + os.path.join(libdir, probe_so) + ): + return libdir + return None + + +def discover_vendor_torch_lib( + bundle_dirname, probe_so, env_override=None, vendor_markers=() +): + """Locate the vendor libtorch .so dir. + + Priority: bundled in this wheel, then ``env_override``, then sibling conda + envs whose torch is a vendor build. + """ + bundled = bundled_lib_dir(bundle_dirname, probe_so) + if bundled: + return bundled + if env_override: + env = os.environ.get(env_override) + if env and os.path.isdir(env): + return env + return _scan_sibling_envs(probe_so, vendor_markers) + + +def _link_one(dst_dir, backup_dir, name, target, required, vendor): + """Idempotently point ``dst_dir/name`` at ``target`` (a vendor .so).""" + dst = os.path.join(dst_dir, name) + if not os.path.exists(target): + if required: + raise FileNotFoundError(f"{vendor} so missing: {target}") + return + # Already correctly linked? + if os.path.islink(dst) and os.path.realpath(dst) == os.path.realpath(target): + return + # Back up a real (non-symlink) original once. + if os.path.exists(dst) and not os.path.islink(dst): + os.makedirs(backup_dir, exist_ok=True) + bak = os.path.join(backup_dir, name) + if not os.path.exists(bak): + os.replace(dst, bak) + else: + os.remove(dst) + elif os.path.islink(dst): + os.remove(dst) # stale/incorrect link + os.symlink(target, dst) + + +def _preload_global(lib_dir, load_order, core_so, vendor, fallback_dir=None): + """dlopen the vendor set RTLD_GLOBAL, in dependency order. + + A CPU-only torch wheel never loads the forked runtime itself. Symlinking the + files is not sufficient on its own: symbols the plugin needs may live in the + forked *CPU* runtime (``GetFlagosDefaultCudaGenerator`` is the measured case + on MetaX) and loading only the vendor library can leave its CPU dependency + RTLD_LOCAL, after which ``libtorch_fl.so`` cannot resolve that symbol. + Loading the whole set globally, core first, avoids that. + + ``lib_dir`` must be the *source* dir (the bundle), never the ``torch/lib`` + symlink dir -- see the ``$ORIGIN`` note in the module docstring. + + ``fallback_dir`` (the stock wheel's ``torch/lib``) covers a non-core .so the + vendor image simply does not ship. Measured: the MetaX CI's + ``/opt/vendor-libtorch/lib`` has no ``libshm.so``, so the bundle has none + either, yet ``libtorch_python.so`` carries a hard ``DT_NEEDED: libshm.so``. + dlopening the stock copy RTLD_GLOBAL *before* ``libtorch_python.so`` + satisfies that DT_NEEDED by soname against the already-loaded object; without + it the loader only searches the bundle's RUNPATH and dies with "libshm.so: + cannot open shared object file". Its own deps (libc10, libtorch_cpu) resolve + through ``torch/lib``, where they are symlinks into the bundle, so they share + an inode with what is already mapped and no second copy appears. + """ + handles = [] + for name in load_order: + path = os.path.join(lib_dir, name) + if not os.path.exists(path): + if name in core_so: + raise FileNotFoundError(f"{vendor} libtorch runtime missing: {path}") + path = os.path.join(fallback_dir, name) if fallback_dir else "" + if not path or not os.path.exists(path): + continue + try: + handles.append(ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL)) + except OSError as exc: + raise RuntimeError( + f"Failed to load {vendor} libtorch runtime: {path}" + ) from exc + return handles + + +def ensure_vendor_libtorch_links( + bundle_dirname, + core_so, + extra_so=(), + env_override=None, + vendor_markers=(), + probe_so=None, + vendor=None, + load_order=None, +): + """Symlink the active torch wheel's core .so to the vendor wheel's copies. + + Args: + bundle_dirname: dir inside the wheel holding the bundle ("lib_maca", ...). + core_so: .so that MUST be present; a missing one raises. + extra_so: .so the stock ``+cpu`` wheel may not ship at all (the vendor + libs); linked in fresh when available, skipped silently otherwise. + env_override: env var naming an explicit source dir. + vendor_markers: substrings identifying a vendor torch in ``version.py``. + probe_so: file whose presence proves a dir is a real vendor libtorch dir. + Defaults to the first entry of ``extra_so``, else of ``core_so``. + vendor: label used in error messages. + load_order: when given, dlopen these RTLD_GLOBAL after linking, in this + order (see ``_preload_global``). Names absent from ``core_so`` may be + missing; a missing core .so raises. + + Returns True if links are in place (or already were), False if there was + nothing to do (no bundle, no vendor torch found, or already running on it). + """ + if bundle_dirname in _done: + return True + probe = probe_so or (extra_so[0] if extra_so else core_so[0]) + label = vendor or bundle_dirname + + active = active_torch_lib() + src = discover_vendor_torch_lib( + bundle_dirname, probe, env_override=env_override, vendor_markers=vendor_markers + ) + if active is None or src is None: + return False + # Already running on the vendor wheel itself -> nothing to do. + if os.path.realpath(active) == os.path.realpath(src): + _done.add(bundle_dirname) + return True + + backup = os.path.join(active, "_orig_backup") + for name in core_so: + _link_one(active, backup, name, os.path.join(src, name), True, label) + for name in extra_so: + _link_one(active, backup, name, os.path.join(src, name), False, label) + + if load_order: + # dlopen from `src` (the bundle), not from `active`: loading through the + # torch/lib symlinks would expand $ORIGIN to torch/lib and lose the + # bundle-internal deps. `active` is only the fallback for a non-core .so + # the vendor image does not ship. Keep the handles alive for the process + # lifetime. + _runtime_handles.extend( + _preload_global(src, load_order, core_so, label, fallback_dir=active) + ) + + _done.add(bundle_dirname) + return True + + +def restore_original_libtorch(core_so, extra_so=(), bundle_dirname=None): + """Undo ensure_vendor_libtorch_links(): drop links, restore the backups.""" + active = active_torch_lib() + if active is None: + return + backup = os.path.join(active, "_orig_backup") + for name in tuple(core_so) + tuple(extra_so): + dst = os.path.join(active, name) + if os.path.islink(dst): + os.remove(dst) + if os.path.isdir(backup): + for name in os.listdir(backup): + os.replace(os.path.join(backup, name), os.path.join(active, name)) + try: + os.rmdir(backup) + except OSError: + pass + if bundle_dirname: + _done.discard(bundle_dirname) diff --git a/torch_fl/accelerator/dcu/__init__.py b/torch_fl/accelerator/dcu/__init__.py new file mode 100644 index 00000000..cc840712 --- /dev/null +++ b/torch_fl/accelerator/dcu/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Hygon DCU (DTK) support for the flagos backend.""" diff --git a/torch_fl/accelerator/dcu/_dcu_libtorch_link.py b/torch_fl/accelerator/dcu/_dcu_libtorch_link.py new file mode 100644 index 00000000..cd64de84 --- /dev/null +++ b/torch_fl/accelerator/dcu/_dcu_libtorch_link.py @@ -0,0 +1,126 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Symlink DTK's libtorch .so into the active (official) torch wheel's lib dir. + +DCU runs the CUDA boxing kernels on top of DTK's hipified libtorch: HIP kernels +are registered under the CUDA dispatch key, so PrivateUse1 -> CUDA re-dispatch +works unchanged. That libtorch is a fork -- measured, DTK's ``libtorch_cpu.so`` +exports 128 hip symbols and carries ``DT_NEEDED: libgalaxyhip.so.5`` -- and it +resolves 2101 of ``libtorch_fl.so``'s undefined symbols, while ``libtorch_hip.so`` +resolves 0. So the core libs are what has to be swapped in; ``libtorch_hip.so`` +still has to be *present* for the dispatch-key registration to exist. + +See ``torch_fl.accelerator._vendor_libtorch`` for the symlink mechanism and why a +ctypes preload cannot replace core libs. + +Unlike MetaX there is no env gate: a DCU wheel is only ever installed on a DCU +box, and the relink is skipped automatically when ``torch_fl/lib_dcu/`` was not +bundled (a plain in-place build) or when torch already IS the DTK wheel. + +The DTK driver stack itself (``libgalaxyhip.so.5``, ``libMIOpen.so.1``, +``librocblas.so.4``, ``librccl.so.1``, ...) stays on the target under +``/opt/dtk`` and is reached via the RUNPATH baked in by +``cmake/FlagosRpath.cmake`` / ``scripts/bundle_dcu_libtorch.sh``. +""" + +from torch_fl.accelerator._vendor_libtorch import ( + bundled_lib_dir, + discover_vendor_torch_lib, + ensure_vendor_libtorch_links, +) +from torch_fl.accelerator._vendor_libtorch import ( + restore_original_libtorch as _restore, +) + +_BUNDLE_DIR = "lib_dcu" + +# Core C++ .so that must come from the DTK wheel as a self-consistent set. +_CORE_SO = ( + "libc10.so", + "libtorch_cpu.so", + "libtorch.so", + "libtorch_global_deps.so", + "libtorch_python.so", +) +# HIP-side .so the stock +cpu wheel does not ship at all. libmagma.so is a +# DT_NEEDED of DTK's libtorch_hip.so; libshm.so is one of libtorch_python.so's +# (torch.multiprocessing's shared-memory manager) and is a *different* build in +# the DTK wheel, so it has to come from the same set. +_HIP_SO = ( + "libc10_hip.so", + "libtorch_hip.so", + "libmagma.so", + "libshm.so", +) + +# Dependency order for the RTLD_GLOBAL preload: core (CPU) first, then the HIP +# side, then libshm and libtorch_python. Same reasoning as MetaX -- symbols the +# plugin needs live in the forked CPU runtime, and loading only the HIP lib can +# leave its CPU dependency RTLD_LOCAL -- including why libshm.so has to be listed +# explicitly ahead of libtorch_python.so. +_LOAD_ORDER = ( + "libc10.so", + "libtorch_cpu.so", + "libtorch.so", + "libtorch_global_deps.so", + "libc10_hip.so", + "libtorch_hip.so", + "libshm.so", + "libtorch_python.so", +) + +_MARKERS = ("dtk", "hip", "das") + + +def _bundled_dcu_lib(): + """Forked libtorch bundled inside this wheel (self-contained DCU build).""" + return bundled_lib_dir(_BUNDLE_DIR, "libtorch_hip.so") + + +def _discover_dcu_torch_lib(): + """Locate the DTK libtorch .so dir. + + Priority: bundled lib_dcu/, then FLAGOS_DCU_TORCH_LIB, then sibling conda + envs whose torch is a DTK build. + """ + return discover_vendor_torch_lib( + _BUNDLE_DIR, + "libtorch_hip.so", + env_override="FLAGOS_DCU_TORCH_LIB", + vendor_markers=_MARKERS, + ) + + +def ensure_dcu_libtorch_links(): + """Symlink the active torch wheel's core .so to DTK's copies. + + Idempotent; reversible via ``torch/lib/_orig_backup/``. Returns True if + links are in place (or already were), False if there was nothing to do. + """ + return ensure_vendor_libtorch_links( + _BUNDLE_DIR, + _CORE_SO, + extra_so=_HIP_SO, + env_override="FLAGOS_DCU_TORCH_LIB", + vendor_markers=_MARKERS, + probe_so="libtorch_hip.so", + vendor="DCU/DTK", + load_order=_LOAD_ORDER, + ) + + +def restore_original_libtorch(): + """Undo ensure_dcu_libtorch_links(): remove links, restore backups.""" + _restore(_CORE_SO, _HIP_SO, bundle_dirname=_BUNDLE_DIR) diff --git a/torch_fl/accelerator/metax/_metax_libtorch_link.py b/torch_fl/accelerator/metax/_metax_libtorch_link.py index d872e27d..78697c66 100644 --- a/torch_fl/accelerator/metax/_metax_libtorch_link.py +++ b/torch_fl/accelerator/metax/_metax_libtorch_link.py @@ -14,34 +14,30 @@ """Symlink MetaX libtorch .so into the active (official) torch wheel's lib dir. -Rationale ---------- On MetaX we reuse PyTorch's CUDA boxing kernels (FLAGOS_METAX_BOXING) by running the *MetaX* C++ runtime (libtorch_cpu.so / libtorch_cuda.so / libc10.so ...), -which is a hard fork exporting ``at::maca::*`` symbols. When the Python front-end -is a *stock* ``torch==X.Y.Z+cpu`` wheel (no CUDA, clean pip env), its own -``torch/lib`` ships the upstream C++ .so. We must make the process load the -MetaX C++ runtime instead. - -Pure ``ctypes`` preloading does NOT reliably work: the official ``_C.so`` / -``libtorch_python.so`` carry an ``$ORIGIN`` RUNPATH that pulls the upstream -``libc10.so`` back in by full path, giving a *second* libc10 in the process and a -duplicate static-init crash (``Key already registered ... caffe2_report_cpu_memory_usage``). - -The robust fix is to make the physical files the RUNPATH resolves to *be* the -MetaX ones -- i.e. replace the stock wheel's ``torch/lib/`` with symlinks to -the MetaX wheel's copies. Originals are backed up to ``torch/lib/_orig_backup/`` -so the operation is fully reversible. - -This runs from ``torch_fl/__init__.py`` BEFORE ``import torch`` (once torch is -imported its libc10 is already mapped and relinking is too late). It is -idempotent, gated on ``FLAGOS_METAX_BOXING=1``, and a no-op when the active torch -already IS the MetaX wheel. +which is a hard fork exporting ``at::maca::*`` symbols. With a stock +``torch==X.Y.Z+cpu`` front-end, the process must load that fork instead of the +upstream .so shipped in ``torch/lib``. + +The mechanism -- symlink replacement, ``_orig_backup/``, the RTLD_GLOBAL preload, +and why a pure ctypes preload cannot do this on its own -- lives in +``torch_fl.accelerator._vendor_libtorch``. This module is just the MetaX .so +lists; whether to relink at all is decided by +``torch_fl.__init__._relink_vendor_libtorch`` (FLAGOS_METAX_BOXING=1 for an +in-place MetaX build, or lib_maca/ being present for a self-contained wheel). """ -import ctypes -import importlib.util -import os +from torch_fl.accelerator._vendor_libtorch import ( + bundled_lib_dir, + discover_vendor_torch_lib, + ensure_vendor_libtorch_links, +) +from torch_fl.accelerator._vendor_libtorch import ( + restore_original_libtorch as _restore, +) + +_BUNDLE_DIR = "lib_maca" # Core C++ .so that must come from the MetaX wheel as a self-consistent set. # libtorch_python.so is included because the stock one references symbols @@ -54,182 +50,87 @@ "libtorch_python.so", ) # CUDA .so the stock +cpu wheel does not ship at all; symlinked in fresh. +# libshm.so is a DT_NEEDED of MetaX's libtorch_python.so (torch.multiprocessing's +# shared-memory manager) and is a *different* build in the MetaX wheel, so it has +# to come from the same set. _CUDA_SO = ( "libc10_cuda.so", "libtorch_cuda.so", "libtorch_cuda_linalg.so", + "libshm.so", ) -_done = False -_runtime_handles = [] - +# Dependency order for the RTLD_GLOBAL preload: core (CPU) first, then the CUDA +# side, then libtorch_python. GetFlagosDefaultCudaGenerator lives in the forked +# ATen CPU runtime, not in libtorch_cuda.so, so loading only the CUDA lib would +# leave that symbol unresolvable from libtorch_fl.so. +# +# libshm.so must precede libtorch_python.so: it is a direct DT_NEEDED of it, and +# dlopen resolves that by soname through the *loader's* search path, not through +# this list. In a fresh install nothing has put libshm.so anywhere the loader +# looks yet, so omitting it here fails with "libshm.so: cannot open shared object +# file" even though the file is sitting in the bundle dir. +_LOAD_ORDER = ( + "libc10.so", + "libtorch_cpu.so", + "libtorch.so", + "libtorch_global_deps.so", + "libc10_cuda.so", + "libtorch_cuda_linalg.so", + "libtorch_cuda.so", + "libshm.so", + "libtorch_python.so", +) -def _active_torch_lib(): - """torch/lib of the importable torch, WITHOUT importing torch.""" - spec = importlib.util.find_spec("torch") - if spec is None or not spec.submodule_search_locations: - return None - lib = os.path.join(spec.submodule_search_locations[0], "lib") - return lib if os.path.isdir(lib) else None +_MARKERS = ("metax", "maca") def _bundled_maca_lib(): - """Forked libtorch bundled inside this wheel (self-contained MetaX build). - - ``scripts/bundle_maca_libtorch.sh`` copies the MetaX libtorch .so into - ``torch_fl/lib_maca/``. When present this is the preferred source: the - target machine then needs only the official ``torch+cpu`` wheel plus the - ``/opt/maca`` driver runtime, no separate MetaX torch wheel. - """ - here = os.path.dirname(os.path.abspath(__file__)) - # this file: torch_fl/accelerator/metax/_metax_libtorch_link.py - pkg_root = os.path.dirname(os.path.dirname(here)) # -> torch_fl/ - libdir = os.path.join(pkg_root, "lib_maca") - if os.path.isdir(libdir) and os.path.exists( - os.path.join(libdir, "libtorch_cuda.so") - ): - return libdir - return None + """Forked libtorch bundled inside this wheel (self-contained MetaX build).""" + return bundled_lib_dir(_BUNDLE_DIR, "libtorch_cuda.so") def _discover_maca_torch_lib(): """Locate the MetaX libtorch .so dir. - Priority: forked libtorch bundled in this wheel (lib_maca/), then an - explicit env var, then sibling conda envs whose torch is a - ``+metax``/``+maca`` build (fallback for multi-env dev setups). + Priority: bundled lib_maca/, then FLAGOS_MACA_TORCH_LIB, then sibling conda + envs whose torch is a ``+metax``/``+maca`` build. """ - bundled = _bundled_maca_lib() - if bundled: - return bundled - - env = os.environ.get("FLAGOS_MACA_TORCH_LIB") - if env and os.path.isdir(env): - return env - - # Scan conda envs next to the current prefix for a MetaX torch build. - prefix = os.environ.get("CONDA_PREFIX") or os.path.dirname( - os.path.dirname(os.__file__) + return discover_vendor_torch_lib( + _BUNDLE_DIR, + "libtorch_cuda.so", + env_override="FLAGOS_MACA_TORCH_LIB", + vendor_markers=_MARKERS, ) - envs_root = os.path.dirname(prefix) # .../envs - if not os.path.isdir(envs_root): - return None - py = "python{}.{}".format(*__import__("sys").version_info[:2]) - for name in sorted(os.listdir(envs_root)): - cand = os.path.join(envs_root, name, "lib", py, "site-packages", "torch") - libdir = os.path.join(cand, "lib") - ver_file = os.path.join(cand, "version.py") - if not os.path.isfile(ver_file) or not os.path.isdir(libdir): - continue - try: - with open(ver_file) as f: - txt = f.read() - except OSError: - continue - if ("metax" in txt or "maca" in txt) and os.path.exists( - os.path.join(libdir, "libtorch_cuda.so") - ): - return libdir - return None - - -def _link_one(dst_dir, backup_dir, name, target, required): - """Idempotently point dst_dir/name at target (a MetaX .so).""" - dst = os.path.join(dst_dir, name) - if not os.path.exists(target): - if required: - raise FileNotFoundError(f"MetaX so missing: {target}") - return - # Already correctly linked? - if os.path.islink(dst) and os.path.realpath(dst) == os.path.realpath(target): - return - # Back up a real (non-symlink) original once. - if os.path.exists(dst) and not os.path.islink(dst): - os.makedirs(backup_dir, exist_ok=True) - bak = os.path.join(backup_dir, name) - if not os.path.exists(bak): - os.replace(dst, bak) - else: - os.remove(dst) - elif os.path.islink(dst): - os.remove(dst) # stale/incorrect link - os.symlink(target, dst) def ensure_maca_libtorch_links(): """Symlink the active torch wheel's core .so to the MetaX wheel's copies. - No-op unless FLAGOS_METAX_BOXING=1. Idempotent; reversible via _orig_backup. - Returns True if links are in place (or already were), False if skipped. + Idempotent; reversible via ``torch/lib/_orig_backup/``. Returns True if + links are in place (or already were), False if there was nothing to do (no + bundle, no MetaX torch found, or torch already IS the MetaX wheel). + + Deciding *whether* to relink is the caller's job -- see + ``torch_fl.__init__._relink_vendor_libtorch``, which gates on + FLAGOS_METAX_BOXING=1 for an in-place MetaX build and on lib_maca/ being + present for a self-contained wheel. This used to self-gate on + FLAGOS_METAX_BOXING, which made the self-contained path a silent no-op: + the stock libtorch_cpu.so stayed in place and libtorch_cuda.so then failed + to resolve at::maca symbols that only the forked CPU runtime defines. """ - global _done, _runtime_handles - if _done: - return True - if os.environ.get("FLAGOS_METAX_BOXING", "0") != "1": - return False - - active = _active_torch_lib() - maca = _discover_maca_torch_lib() - if active is None or maca is None: - return False - # Already running on the MetaX wheel itself -> nothing to do. - if os.path.realpath(active) == os.path.realpath(maca): - _done = True - return True - - backup = os.path.join(active, "_orig_backup") - for name in _CORE_SO: - _link_one(active, backup, name, os.path.join(maca, name), required=True) - for name in _CUDA_SO: - _link_one(active, backup, name, os.path.join(maca, name), required=False) - - # A CPU-only torch wheel does not load the forked runtime itself. Load the - # complete bundled dependency set globally before torch_fl._C loads. In - # particular, GetFlagosDefaultCudaGenerator is provided by the forked - # ATen runtime and is not guaranteed to be in libtorch_cuda.so itself. - # Loading only the CUDA library can leave its CPU dependency RTLD_LOCAL, - # so libtorch_fl.so cannot resolve that symbol later. - load_order = ( - "libc10.so", - "libtorch_cpu.so", - "libtorch.so", - "libtorch_global_deps.so", - "libc10_cuda.so", - "libtorch_cuda_linalg.so", - "libtorch_cuda.so", - "libtorch_python.so", + return ensure_vendor_libtorch_links( + _BUNDLE_DIR, + _CORE_SO, + extra_so=_CUDA_SO, + env_override="FLAGOS_MACA_TORCH_LIB", + vendor_markers=_MARKERS, + probe_so="libtorch_cuda.so", + vendor="MetaX", + load_order=_LOAD_ORDER, ) - for name in load_order: - path = os.path.join(active, name) - if not os.path.exists(path): - if name in _CORE_SO: - raise FileNotFoundError(f"MetaX libtorch runtime missing: {path}") - continue - try: - _runtime_handles.append(ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL)) - except OSError as exc: - raise RuntimeError( - f"Failed to load MetaX libtorch runtime: {path}" - ) from exc - - _done = True - return True def restore_original_libtorch(): """Undo ensure_maca_libtorch_links(): remove links, restore backups.""" - active = _active_torch_lib() - if active is None: - return - backup = os.path.join(active, "_orig_backup") - for name in _CORE_SO + _CUDA_SO: - dst = os.path.join(active, name) - if os.path.islink(dst): - os.remove(dst) - if os.path.isdir(backup): - for name in os.listdir(backup): - os.replace(os.path.join(backup, name), os.path.join(active, name)) - try: - os.rmdir(backup) - except OSError: - pass + _restore(_CORE_SO, _CUDA_SO, bundle_dirname=_BUNDLE_DIR) diff --git a/torch_fl/accelerator/ppu/__init__.py b/torch_fl/accelerator/ppu/__init__.py new file mode 100644 index 00000000..da8ad8f7 --- /dev/null +++ b/torch_fl/accelerator/ppu/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""T-Head PPU support for the flagos backend.""" diff --git a/torch_fl/accelerator/ppu/_ppu_libtorch_link.py b/torch_fl/accelerator/ppu/_ppu_libtorch_link.py new file mode 100644 index 00000000..09e4e103 --- /dev/null +++ b/torch_fl/accelerator/ppu/_ppu_libtorch_link.py @@ -0,0 +1,121 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Symlink the PPU libtorch .so into the active (official) torch wheel's lib dir. + +PPU builds against ``PPU_SDK/CUDA_SDK``, so ``ACCELERATOR=cuda`` and the CUDA +boxing kernels apply unchanged. What differs from a real NVIDIA box is the +libtorch: it is a local ``USE_CUDA=1`` source build, not the upstream wheel, and +it resolves 2092 of ``libtorch_fl.so``'s undefined symbols (``libtorch_cuda.so`` +resolves 0, ``libc10_cuda.so`` 10) -- so the core libs must be swapped in, and +``libtorch_cuda.so`` must merely be present for the CUDA dispatch keys. + +That local build also links the system MKL from ``/usr/local/lib`` +(``libmkl_core``/``libmkl_gnu_thread``/``libmkl_intel_lp64``), which the bundling +script copies into ``lib_ppu/`` alongside the core libs. + +See ``torch_fl.accelerator._vendor_libtorch`` for the symlink mechanism and why a +ctypes preload cannot replace core libs. No env gate: the relink is skipped when +``torch_fl/lib_ppu/`` was not bundled or torch already IS the PPU build. The PPU +SDK runtime stays on the target under ``/usr/local/PPU_SDK``. +""" + +from torch_fl.accelerator._vendor_libtorch import ( + bundled_lib_dir, + discover_vendor_torch_lib, + ensure_vendor_libtorch_links, +) +from torch_fl.accelerator._vendor_libtorch import ( + restore_original_libtorch as _restore, +) + +_BUNDLE_DIR = "lib_ppu" + +# Core C++ .so that must come from the PPU build as a self-consistent set. +_CORE_SO = ( + "libc10.so", + "libtorch_cpu.so", + "libtorch.so", + "libtorch_global_deps.so", + "libtorch_python.so", +) +# CUDA-side .so the stock +cpu wheel does not ship at all. +# libshm.so is a DT_NEEDED of PPU's libtorch_python.so (torch.multiprocessing's +# shared-memory manager) and is a *different* build in the PPU wheel, so it has +# to come from the same set. +_CUDA_SO = ( + "libc10_cuda.so", + "libtorch_cuda.so", + "libtorch_cuda_linalg.so", + "libshm.so", +) + +# Dependency order for the RTLD_GLOBAL preload: core (CPU) first, then the CUDA +# side, then libshm and libtorch_python. Same reasoning as MetaX, including why +# libshm.so has to be listed explicitly ahead of libtorch_python.so. +_LOAD_ORDER = ( + "libc10.so", + "libtorch_cpu.so", + "libtorch.so", + "libtorch_global_deps.so", + "libc10_cuda.so", + "libtorch_cuda_linalg.so", + "libtorch_cuda.so", + "libshm.so", + "libtorch_python.so", +) + +_MARKERS = ("ppu",) + + +def _bundled_ppu_lib(): + """PPU libtorch bundled inside this wheel (self-contained PPU build).""" + return bundled_lib_dir(_BUNDLE_DIR, "libtorch_cuda.so") + + +def _discover_ppu_torch_lib(): + """Locate the PPU libtorch .so dir. + + Priority: bundled lib_ppu/, then FLAGOS_PPU_TORCH_LIB, then sibling conda + envs whose torch is a PPU build. + """ + return discover_vendor_torch_lib( + _BUNDLE_DIR, + "libtorch_cuda.so", + env_override="FLAGOS_PPU_TORCH_LIB", + vendor_markers=_MARKERS, + ) + + +def ensure_ppu_libtorch_links(): + """Symlink the active torch wheel's core .so to the PPU build's copies. + + Idempotent; reversible via ``torch/lib/_orig_backup/``. Returns True if + links are in place (or already were), False if there was nothing to do. + """ + return ensure_vendor_libtorch_links( + _BUNDLE_DIR, + _CORE_SO, + extra_so=_CUDA_SO, + env_override="FLAGOS_PPU_TORCH_LIB", + vendor_markers=_MARKERS, + probe_so="libtorch_cuda.so", + vendor="PPU", + load_order=_LOAD_ORDER, + ) + + +def restore_original_libtorch(): + """Undo ensure_ppu_libtorch_links(): remove links, restore backups.""" + _restore(_CORE_SO, _CUDA_SO, bundle_dirname=_BUNDLE_DIR) diff --git a/torch_fl/csrc/CMakeLists.txt b/torch_fl/csrc/CMakeLists.txt index 6e7fa3e5..46993560 100644 --- a/torch_fl/csrc/CMakeLists.txt +++ b/torch_fl/csrc/CMakeLists.txt @@ -36,6 +36,12 @@ endif() target_link_directories(${LIBRARY_NAME} PRIVATE ${PYTORCH_INSTALL_DIR}/lib) +# Same reasoning as csrc/CMakeLists.txt: keep the build machine's torch/lib out +# of the installed RUNPATH and resolve a bundled forked libtorch via $ORIGIN. +# libtorch_bindings.so is loaded as torch_fl._C, i.e. always after +# torch_fl/__init__.py has imported torch. +flagos_set_portable_rpath(${LIBRARY_NAME}) + install(TARGETS ${LIBRARY_NAME} ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}