From 8a69ba3d72b5ec6c5fc078896d708d4c36bbc83c Mon Sep 17 00:00:00 2001 From: lvyufeng Date: Fri, 7 Aug 2026 23:16:00 +0800 Subject: [PATCH] feat(bpu): D-Robotics RDK BPU backend with on-board compilation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge the RDK BPU support into torch-fl as ACCELERATOR=bpu. This is a graph-compile backend (torch.compile(backend="bpu")) rather than per-op kernels: the BPU executes whole .hbm artifacts produced by hbdk4, so there are no PrivateUse1 operator registrations and eager ops reach cpu_fallback. Key decisions: - Eager: real device memory (hbUCPMallocCached) + CPU compute (no per-op BPU kernels). Tensors genuinely live in UCP memory so the compile path can be zero-copy where dtype matches. - Compile: on-board via box64-wrapped x86_64 hbdk4. No VM, no separate host, and no kernel rebuild — the earlier conclusion that a 4KB-page kernel was required turned out to be wrong; current box64 handles 64KB pages itself. - One torch-fl package: no separate torch_bpu repo. ACCELERATOR=bpu builds it, import torch_fl registers the backend. Build system: - CMakeLists.txt: bpu branch bypasses CUDA requirements, links Horizon runtime (libbpu, libhbucp, libhbdnn). - setup.py: bpu turns off all kernel sets (CUDA/FlagGems/MetaX/Ascend), installs torch_fl/accelerator/bpu/*.py. - csrc/runtime/accelerator/bpu/: UCP allocator + device/stream/event stubs. The Python package sits in torch_fl/accelerator/ alongside cuda/dcu/gcu/ metax/ppu, not in a new torch_fl/backends/. "backends" is already taken in this tree: torch_fl/backends.conf and configs/backends_*.conf are the op routing tables, and torch.backends upstream means feature flags. PyTorch 2.10 pin: - pyproject.toml + setup.py: require torch>=2.10,<2.11. The checked-in csrc/aten/generated/* is version-sensitive; a mismatch surfaces as compile errors, not resolver failures. Convolution registration (csrc/aten/register.cc): - aten::convolution dispatches PrivateUse1 to convolution_overrideable, and the only other kernel is a CompositeExplicitAutograd stub that raises. The boxed cpu_fallback cannot help: it moves args to CPU and redispatches the *same* op, landing back on the stub. Register explicit wrappers that call at::convolution_symint on CPU tensors. On-board hbdk4 (torch_fl/accelerator/bpu/compiler.py): - x86_emulator() probes box64/qemu-x86_64-static with a source-built box64 as FLAGOS_BPU_X86_EMULATOR. Distro box64 0.2.6 fails; current versions handle 64KB pages at runtime. - x86_env() sets BOX64_LD_LIBRARY_PATH (_mlir_libs for the hbdk4 .so tree), PYTHONPATH (stubs for numba/torch — both cause segfaults under box64 but are only imported, never executed on the ONNX path), and LD_PRELOAD (libhbtl.so with RTLD_GLOBAL — _hbdk.so needs hbtl symbols but doesn't list libhbtl in DT_NEEDED). - _compile_via_x86_emulator() tolerates post-compile AllocError (hbdk4 loads the artifact back for validation, which can fail in emulation after a complete .hbm is written). - scripts/setup_bpu_hbdk4.sh: idempotent setup (box64 build, x86 python, hbdk4 wheels, stubs). Zero-copy (torch_fl/accelerator/bpu/runtime.py): - _device_view(): wrap a flagos tensor's UCP storage in a numpy array in-place (verified: address == data_ptr(), writes through the view are visible to torch). No D2H copy inbound. Quantization still copies (dtype changes float32->int8), and hbm_runtime.run allocates its own outputs. Exit-time leak (csrc/runtime/allocator/caching_device_allocator.cc): - Skip block release on BPU, same as tsingmicro. The allocator is a function-local static, so its dtor runs from __run_exit_handlers after libhbucp's FINI_ARRAY may have released the heap. Calling torch_fl._C._empty_cache() before exit frees cleanly, confirming only the ordering is at fault. Tests (tests/unit/bpu/): - 8 files, 65 tests covering partition, qdq, decompose, freeze, splice, eager device ops (elementwise, matmul, reduction, conv2d forward/backward), zero-copy numpy views, and x86 environment construction. - All BPU tests pass (65); the 4 profiler failures on this board predate this branch (they need libcupti and real GPU kernels). ruff clean. - No shared code touched: csrc/aten/generated/* and the non-BPU backends_*.conf are byte-identical to main. Measured performance: 6-layer conv stack @224x224, 72.06ms eager -> 3.75ms BPU (19.2x), int8 quantized with ~3% relative error. Docs: docs/bpu.md covers architecture, quantization, weight freezing, on-board hbdk4 setup, calibration, environment variables, and known limits. Co-Authored-By: Claude Opus 5 (1M context) --- CMakeLists.txt | 33 +- README.md | 60 ++- csrc/CMakeLists.txt | 24 +- csrc/aten/common.cc | 2 + csrc/aten/contiguous_ops.cc | 4 +- csrc/aten/copy_ops.cc | 13 +- csrc/aten/register.cc | 99 ++++ csrc/runtime/accelerator/CMakeLists.txt | 27 + csrc/runtime/accelerator/bpu/bpu.h | 38 ++ csrc/runtime/accelerator/bpu/device.cc | 84 +++ csrc/runtime/accelerator/bpu/memory.cc | 269 ++++++++++ csrc/runtime/accelerator/bpu/stream.cc | 145 +++++ csrc/runtime/allocator/backends/bpu_memory.h | 79 +++ .../allocator/caching_device_allocator.cc | 21 +- csrc/runtime/guard.h | 2 +- csrc/runtime/hooks.h | 5 +- docs/bpu.md | 281 ++++++++++ pyproject.toml | 5 +- scripts/setup_bpu_hbdk4.sh | 298 +++++++++++ setup.py | 30 +- tests/unit/bpu/test_decompose.py | 119 +++++ tests/unit/bpu/test_eager_device.py | 137 +++++ tests/unit/bpu/test_freeze.py | 182 +++++++ tests/unit/bpu/test_partition.py | 170 ++++++ tests/unit/bpu/test_qdq.py | 222 ++++++++ tests/unit/bpu/test_splice.py | 198 +++++++ tests/unit/bpu/test_x86_env.py | 142 +++++ tests/unit/bpu/test_zero_copy.py | 108 ++++ torch_fl/__init__.py | 35 ++ torch_fl/accelerator/bpu/__init__.py | 36 ++ torch_fl/accelerator/bpu/backend.py | 287 ++++++++++ torch_fl/accelerator/bpu/calibrate.py | 224 ++++++++ torch_fl/accelerator/bpu/compiler.py | 495 ++++++++++++++++++ torch_fl/accelerator/bpu/decompose.py | 100 ++++ torch_fl/accelerator/bpu/partition.py | 307 +++++++++++ torch_fl/accelerator/bpu/qdq.py | 142 +++++ torch_fl/accelerator/bpu/runtime.py | 265 ++++++++++ torch_fl/configs/backends_bpu.conf | 17 + 38 files changed, 4682 insertions(+), 23 deletions(-) create mode 100644 csrc/runtime/accelerator/bpu/bpu.h create mode 100644 csrc/runtime/accelerator/bpu/device.cc create mode 100644 csrc/runtime/accelerator/bpu/memory.cc create mode 100644 csrc/runtime/accelerator/bpu/stream.cc create mode 100644 csrc/runtime/allocator/backends/bpu_memory.h create mode 100644 docs/bpu.md create mode 100755 scripts/setup_bpu_hbdk4.sh create mode 100644 tests/unit/bpu/test_decompose.py create mode 100644 tests/unit/bpu/test_eager_device.py create mode 100644 tests/unit/bpu/test_freeze.py create mode 100644 tests/unit/bpu/test_partition.py create mode 100644 tests/unit/bpu/test_qdq.py create mode 100644 tests/unit/bpu/test_splice.py create mode 100644 tests/unit/bpu/test_x86_env.py create mode 100644 tests/unit/bpu/test_zero_copy.py create mode 100644 torch_fl/accelerator/bpu/__init__.py create mode 100644 torch_fl/accelerator/bpu/backend.py create mode 100644 torch_fl/accelerator/bpu/calibrate.py create mode 100644 torch_fl/accelerator/bpu/compiler.py create mode 100644 torch_fl/accelerator/bpu/decompose.py create mode 100644 torch_fl/accelerator/bpu/partition.py create mode 100644 torch_fl/accelerator/bpu/qdq.py create mode 100644 torch_fl/accelerator/bpu/runtime.py create mode 100644 torch_fl/configs/backends_bpu.conf diff --git a/CMakeLists.txt b/CMakeLists.txt index 805feceb..5feabb62 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -14,7 +14,7 @@ cmake_minimum_required(VERSION 3.18 FATAL_ERROR) -set(ACCELERATOR "cuda" CACHE STRING "Accelerator platform: cuda, metax, ascend, tsingmicro, dcu, gcu, or musa") +set(ACCELERATOR "cuda" CACHE STRING "Accelerator platform: cuda, metax, ascend, tsingmicro, dcu, gcu, musa, or bpu") # Directory inside the wheel holding a bundled forked libtorch, when the backend # ships one (see scripts/bundle_*_libtorch.sh). "lib" means "no separate bundle @@ -95,6 +95,14 @@ elseif(ACCELERATOR STREQUAL "musa") set(CUDA_KERNEL OFF CACHE BOOL "Build CUDA kernel implementations" FORCE) set(FLAGGEMS_KERNEL OFF CACHE BOOL "Build FlagGems kernel implementations" FORCE) project(TORCH_FLAGOS CXX C) +elseif(ACCELERATOR STREQUAL "bpu") + # D-Robotics RDK BPU. No CUDA runtime and no per-op kernels: the BPU executes + # whole compiled graphs, so eager ops fall back to CPU and acceleration comes + # from the torch.compile backend. Only the runtime layer (UCP allocator, + # device/stream stubs) is native, all plain host C++. + set(CUDA_KERNEL OFF CACHE BOOL "Build CUDA kernel implementations" FORCE) + 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 @@ -133,7 +141,8 @@ set(CMAKE_CXX_EXTENSIONS OFF) if(NOT ACCELERATOR STREQUAL "metax" AND NOT ACCELERATOR STREQUAL "ascend" AND NOT ACCELERATOR STREQUAL "tsingmicro" AND NOT ACCELERATOR STREQUAL "dcu" - AND NOT ACCELERATOR STREQUAL "gcu" AND NOT ACCELERATOR STREQUAL "musa") + AND NOT ACCELERATOR STREQUAL "gcu" AND NOT ACCELERATOR STREQUAL "musa" + AND NOT ACCELERATOR STREQUAL "bpu") set(CMAKE_CUDA_STANDARD 17) set(CMAKE_CUDA_STANDARD_REQUIRED ON) endif() @@ -376,6 +385,26 @@ elseif(ACCELERATOR STREQUAL "musa") message(STATUS "mudnn op library: ${MUDNN_LIB}") endif() + add_library(torch::cudart INTERFACE IMPORTED) + if(NOT TARGET CUDA::nvToolsExt) + add_library(CUDA::nvToolsExt INTERFACE IMPORTED) + endif() +elseif(ACCELERATOR STREQUAL "bpu") + # D-Robotics RDK BPU. There is no CUDA-compat layer here at all: device + # memory comes from the Horizon UCP allocator (libhbucp) and the core/task API + # from libbpu, both linked by csrc/runtime/accelerator/bpu directly. So there + # is no CUDA_RUNTIME_LIB to hand to the other targets. + set(HOBOT_LIB_DIR "/usr/hobot/lib" CACHE PATH "Horizon runtime library dir") + set(HOBOT_INCLUDE_DIR "/usr/include/hobot" CACHE PATH "Horizon runtime include dir") + message(STATUS "Horizon runtime: ${HOBOT_LIB_DIR} (headers ${HOBOT_INCLUDE_DIR})") + + set(CUDA_RUNTIME_LIB "") + + # Embed the Horizon library path so LD_LIBRARY_PATH is not needed at runtime. + list(APPEND CMAKE_INSTALL_RPATH "${HOBOT_LIB_DIR}") + + # Pre-create torch::cudart so PyTorch's cuda.cmake returns early + # (guard: if(TARGET torch::cudart) return()). add_library(torch::cudart INTERFACE IMPORTED) if(NOT TARGET CUDA::nvToolsExt) add_library(CUDA::nvToolsExt INTERFACE IMPORTED) diff --git a/README.md b/README.md index 03c2eb25..a3d85401 100644 --- a/README.md +++ b/README.md @@ -609,11 +609,62 @@ then rejects. - Build with `-DMUSA_KERNEL=OFF` to skip mudnn entirely; the runtime still works and all compute falls back to CPU. +### Build from Source (D-Robotics RDK BPU Platform) + +The BPU is the one platform here with **no operator kernels at all**. Its BPU +executes whole compiled graphs (a `.hbm` produced by hbdk4), not individual ops, +so there is nothing for a `PrivateUse1` kernel to call. `torch_fl` supplies a +real device — UCP-backed memory plus the device/stream layer — and a +`torch.compile` backend; eager ops run on the CPU. + +The supported target is **`torch==2.10.0+cpu`** (the cp314 aarch64 wheel on +PyPI is what the board runs), and `setup.py`'s `TORCH_PIN` enforces +`torch>=2.10,<2.11` here as on every other platform. That pin is not +incidental: the checked-in `csrc/aten/generated/*` bindings are generated +against one ATen surface, so a newer torch fails as a wall of compile errors +rather than a clean resolver error. + +```bash +# Upstream CPU torch wheel; the Horizon runtime ships in the board image. +pip install torch==2.10.0+cpu --index-url https://download.pytorch.org/whl/cpu +ACCELERATOR=bpu pip install --no-build-isolation -vvv -e . +``` + +```python +import torch, torch_fl + +compiled = torch.compile(MyNet().eval(), backend="bpu") +out = compiled(torch.randn(1, 3, 224, 224)) +``` + +**BPU-specific notes:** + +- **No SDK root to configure.** `libhbucp.so` (UCP allocator) and `libbpu.so` + (core/task API) ship at `/usr/hobot/lib` with headers at `/usr/include/hobot`. +- **hbdk4 compiles on the board, on the stock kernel.** The compiler ships + x86_64-only wheels, so it runs under box64 — which needs to be **built from + source**: the packaged 0.2.6 aborts on this board's 64 KB pages, while 0.4+ + handles them at runtime. No VM, no cross-compile host, no kernel rebuild. + Run `scripts/setup_bpu_hbdk4.sh` once, then export + `FLAGOS_BPU_X86_PYTHON` and `FLAGOS_BPU_X86_EMULATOR`. Without a reachable + hbdk4 the backend logs a warning and runs every partition on the CPU, so the + install is still usable. Details in [docs/bpu.md](docs/bpu.md). +- **Quantization is a precondition, not an optimization.** hbdk4 lowers float + conv to the CPU, so int8 Q/DQ insertion is what puts work on the BPU at all. + On by default; `FLAGOS_BPU_QUANTIZE=0` for bit-exact float artifacts. +- **Convolution is registered explicitly.** `aten::convolution` routes + `PrivateUse1` to `convolution_overrideable`, whose only other kernel raises, + so the boxed fallback cannot reach a CPU implementation. Two wrappers in + `register.cc` call `at::convolution` on CPU tensors instead. +- Measured on-board (torch 2.10): a 6-layer conv stack at 224x224 runs **3.75 ms + on the BPU vs 72.06 ms eager CPU — 19.2x**. Toy nets are a wash — submission + overhead dominates. + ### Build Environment Variables | Variable | Description | |----------|-------------| -| `ACCELERATOR` | Hardware platform: `cuda` (default), `metax`, `ascend`, `tsingmicro`, `dcu`, `gcu`, or `musa` | +| `ACCELERATOR` | Hardware platform: `cuda` (default), `metax`, `ascend`, `tsingmicro`, `dcu`, `gcu`, `musa`, or `bpu` | | `CUDA_HOME` | CUDA toolkit path | | `DTK_ROOT` | Hygon DTK path (falls back to `ROCM_PATH`, then `/opt/dtk`; required for DCU build) | | `TOPS_HOME` | Enflame TopsRider SDK path (default `/opt/tops`; required for GCU build) | @@ -629,6 +680,9 @@ then rejects. | `GCU_KERNEL` | Enable Enflame GCU topsaten kernel build (`ON`/`OFF`, auto-enabled when `ACCELERATOR=gcu`) | | `MUSA_HOME` | Moore Threads MUSA toolkit path (default `/usr/local/musa`; required for MUSA build) | | `MUSA_KERNEL` | Enable the MUSA mudnn kernel build (`ON`/`OFF`, auto-enabled when `ACCELERATOR=musa`); `OFF` falls back to CPU for all compute | +| `FLAGOS_BPU_X86_PYTHON` | Path to an x86_64 python with `hbdk4-compiler`, run under box64 (BPU compile path; see [docs/bpu.md](docs/bpu.md)) | +| `FLAGOS_BPU_X86_EMULATOR` | Path to a box64 binary (0.4+); the packaged 0.2.6 cannot run on this board's 64 KB pages | +| `FLAGOS_BPU_X86_STUBS` | numba/torch import stubs for the emulated interpreter (defaults next to the x86 python) | ### Runtime Environment Variables @@ -644,6 +698,10 @@ then rejects. | `FLAGOS_LOG_DISPATCH` | Set to `1` to print backend selection for each operator dispatch | | `FLAGOS_OP_` | Per-operator backend override (replace `.` with `__` in op names) | | `TORCH_DEVICE_BACKEND_AUTOLOAD` | Set to `0` to stop a vendor plugin (e.g. `torch_musa`) from claiming `PrivateUse1` during `import torch`; `torch_fl` sets this itself on MUSA builds | +| `FLAGOS_BPU_MARCH` | BPU micro-architecture (default `nash-p`) | +| `FLAGOS_BPU_QUANTIZE` | BPU int8 Q/DQ insertion (default `1`; `0` compiles float, which keeps conv on the CPU) | +| `FLAGOS_BPU_ACT_SCALE` | BPU fallback activation scale for uncalibrated tensors (default `0.05`) | +| `FLAGOS_BPU_CACHE` | BPU `.hbm` artifact cache (default `~/.cache/torch_fl_bpu`) | ## Usage diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index af5c4f01..c0f0abb0 100644 --- a/csrc/CMakeLists.txt +++ b/csrc/CMakeLists.txt @@ -290,6 +290,19 @@ elseif(ACCELERATOR STREQUAL "musa") # vendor's torch version). target_link_libraries(${LIBRARY_NAME} PRIVATE ${MUDNN_LIB}) endif() +elseif(ACCELERATOR STREQUAL "bpu") + # Selects backends/bpu_memory.h and, like tsingmicro, skips the generated + # PrivateUse1 op registrations: the BPU executes whole compiled graphs, so + # there is no per-op kernel to register and every aten call reaches + # cpu_fallback. Acceleration comes from the torch.compile backend instead. + target_compile_definitions(${LIBRARY_NAME} PRIVATE USE_BPU=1) + target_link_libraries(${LIBRARY_NAME} PRIVATE ${_torch_fl_link_libs}) + find_path(HOBOT_UCP_INCLUDE_DIR hb_ucp_sys.h + PATHS ${HOBOT_INCLUDE_DIR} /usr/include/hobot) + if(HOBOT_UCP_INCLUDE_DIR) + target_include_directories(${LIBRARY_NAME} PRIVATE + ${HOBOT_UCP_INCLUDE_DIR} ${HOBOT_UCP_INCLUDE_DIR}/dnn) + endif() elseif(ACCELERATOR STREQUAL "gcu") target_compile_definitions(${LIBRARY_NAME} PRIVATE USE_GCU=1) target_link_libraries(${LIBRARY_NAME} PRIVATE ${_torch_fl_link_libs}) @@ -343,10 +356,13 @@ install(TARGETS ${LIBRARY_NAME} RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR} ) -if(ACCELERATOR STREQUAL "gcu" OR ACCELERATOR STREQUAL "musa") - # torch_fl/__init__.py pins backends_cuda.conf by default, which on GCU/MUSA - # would route every op to the (absent) CUDA kernels. Record the platform next - # to the libs so the Python side can pick backends_.conf instead. +if(ACCELERATOR STREQUAL "gcu" OR ACCELERATOR STREQUAL "musa" + OR ACCELERATOR STREQUAL "bpu") + # torch_fl/__init__.py pins backends_cuda.conf by default, which on + # GCU/MUSA/BPU would route every op to the (absent) CUDA kernels. Record the + # platform next to the libs so the Python side can pick + # backends_.conf instead. BPU additionally uses the marker to decide + # whether to register the "bpu" torch.compile backend. set(_flagos_platform_marker "${CMAKE_CURRENT_BINARY_DIR}/flagos_platform") file(WRITE "${_flagos_platform_marker}" "${ACCELERATOR}\n") install(FILES "${_flagos_platform_marker}" diff --git a/csrc/aten/common.cc b/csrc/aten/common.cc index 958e8ccd..8a675bb7 100644 --- a/csrc/aten/common.cc +++ b/csrc/aten/common.cc @@ -37,6 +37,8 @@ std::string DefaultConfigPath() { platform = "ascend"; #elif defined(USE_MUSA) platform = "musa"; +#elif defined(USE_BPU) + platform = "bpu"; #endif if (platform) { // dir is /torch_fl/lib, configs are at /torch_fl/configs/ diff --git a/csrc/aten/contiguous_ops.cc b/csrc/aten/contiguous_ops.cc index e6155b36..0c9120db 100644 --- a/csrc/aten/contiguous_ops.cc +++ b/csrc/aten/contiguous_ops.cc @@ -39,7 +39,7 @@ at::Tensor contiguous( // intermediate buffer and no CPU round-trip. musa_ops::MudnnCopy(self, result); #elif !defined(USE_ASCEND) && !defined(USE_TSINGMICRO) && !defined(USE_GCU) && \ - !defined(USE_MUSA) + !defined(USE_MUSA) && !defined(USE_BPU) // CUDA platform: use DeviceBoxingGuard to invoke native CUDA strided copy // kernel on-device, avoiding expensive CPU round-trip. DeviceBoxingGuard guard(self, result); @@ -105,7 +105,7 @@ at::Tensor clone( // strided copy kernel instead of expensive CPU round-trip. if (self.is_privateuseone()) { #if !defined(USE_ASCEND) && !defined(USE_TSINGMICRO) && !defined(USE_GCU) && \ - !defined(USE_MUSA) + !defined(USE_MUSA) && !defined(USE_BPU) auto result = at::empty( self.sizes(), self.options().memory_format(memory_format)); DeviceBoxingGuard guard(self, result); diff --git a/csrc/aten/copy_ops.cc b/csrc/aten/copy_ops.cc index 1f8f7778..9c864855 100644 --- a/csrc/aten/copy_ops.cc +++ b/csrc/aten/copy_ops.cc @@ -35,7 +35,7 @@ // 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_DCU) + !defined(USE_MUSA) && !defined(USE_DCU) && !defined(USE_BPU) #define FLAGOS_COPY_HAS_CUDA_STREAM 1 #include #endif @@ -132,7 +132,7 @@ at::Tensor _copy_from( // this platform. musa_ops::MudnnCopy(self, const_cast(dst)); #elif !defined(USE_ASCEND) && !defined(USE_TSINGMICRO) && !defined(USE_GCU) && \ - !defined(USE_MUSA) + !defined(USE_MUSA) && !defined(USE_BPU) // CUDA platform: use DeviceBoxingGuard to dispatch to native CUDA // strided copy kernel (handles strides, dtype casts on-device). DeviceBoxingGuard guard(self, dst); @@ -214,7 +214,7 @@ at::Tensor _copy_from( auto tmp = at::empty(self_contig.sizes(), dst.options()); Memcpy(tmp.data_ptr(), self_contig.data_ptr(), nbytes, MemcpyHostToDevice); #if defined(USE_ASCEND) || defined(USE_TSINGMICRO) || defined(USE_GCU) || \ - defined(USE_MUSA) + defined(USE_MUSA) || defined(USE_BPU) at::native::flagos::_copy_from(tmp, dst, false); #else DeviceBoxingGuard guard(tmp, dst); @@ -244,7 +244,7 @@ at::Tensor _copy_from( auto tmp = at::empty(self_contig.sizes(), dst.options()); Memcpy(tmp.data_ptr(), self_contig.data_ptr(), nbytes, MemcpyDeviceToDevice); #if defined(USE_ASCEND) || defined(USE_TSINGMICRO) || defined(USE_GCU) || \ - defined(USE_MUSA) + defined(USE_MUSA) || defined(USE_BPU) at::native::flagos::_copy_from(tmp, dst, false); #else DeviceBoxingGuard guard(tmp, dst); @@ -405,7 +405,7 @@ at::Tensor _to_copy( .dtype(dtype).device(c10::Device(c10::kPrivateUse1, device_index))); musa_ops::MudnnCopy(self_contig, result); #elif defined(USE_ASCEND) || defined(USE_TSINGMICRO) || defined(USE_GCU) || \ - defined(USE_MUSA) + defined(USE_MUSA) || defined(USE_BPU) // No CUDA runtime on these backends, so the CUDA TensorIterator cast // below is unavailable. #ifdef USE_ASCEND @@ -415,7 +415,8 @@ at::Tensor _to_copy( #endif if (!result.defined()) { // Fallback: CPU round-trip when no on-device cast is available - // (TsingMicro / GCU / MUSA, or an Ascend dtype pair aclnnCast rejects). + // (TsingMicro / GCU / MUSA / BPU, or an Ascend dtype pair aclnnCast + // rejects). size_t nbytes = self_contig.numel() * self_contig.element_size(); at::Tensor cpu_tensor = at::empty(self_contig.sizes(), self_contig.options().device(at::kCPU)); diff --git a/csrc/aten/register.cc b/csrc/aten/register.cc index 84bbfb1d..6886bde7 100644 --- a/csrc/aten/register.cc +++ b/csrc/aten/register.cc @@ -230,6 +230,86 @@ int64_t WrapperFusedSdpChoice( return static_cast(at::SDPBackend::efficient_attention); } +#if defined(USE_BPU) +// Convolution on a device with no per-op kernels. +// +// aten::convolution dispatches PrivateUse1 to convolution_overrideable, and the +// only other kernel registered for that op is a CompositeExplicitAutograd stub +// that raises NotImplementedError. So the boxed cpu_fallback cannot help here: +// it moves the arguments to CPU and redispatches the same op, which lands back +// on the stub. These wrappers cross to CPU and then call at::convolution -- +// a different op, and the one that actually has a CPU kernel. +at::Tensor BPUWrapperConvolutionOverrideable( + const at::Tensor& input, + const at::Tensor& weight, + const ::std::optional& bias, + c10::SymIntArrayRef stride, + c10::SymIntArrayRef padding, + c10::SymIntArrayRef dilation, + bool transposed, + c10::SymIntArrayRef output_padding, + c10::SymInt groups) { + auto out = at::convolution_symint( + input.cpu(), + weight.cpu(), + bias.has_value() && bias->defined() + ? ::std::optional(bias->cpu()) + : ::std::nullopt, + stride, + padding, + dilation, + transposed, + output_padding, + groups); + return out.to(input.device()); +} + +::std::tuple +BPUWrapperConvolutionBackwardOverrideable( + const at::Tensor& grad_output, + const at::Tensor& input, + const at::Tensor& weight, + c10::SymIntArrayRef stride, + c10::SymIntArrayRef padding, + c10::SymIntArrayRef dilation, + bool transposed, + c10::SymIntArrayRef output_padding, + c10::SymInt groups, + ::std::array output_mask) { + // convolution_backward wants the forward bias *sizes*, not the bias, and only + // to shape grad_bias -- which is a plain sum over the non-channel dims, so the + // channel count is all it needs. + ::std::optional bias_sizes = ::std::nullopt; + c10::SymInt out_channels = + transposed ? weight.sym_size(1) * groups : weight.sym_size(0); + ::std::vector bias_shape{out_channels}; + if (output_mask[2]) { + bias_sizes = c10::SymIntArrayRef(bias_shape); + } + + auto out = at::convolution_backward_symint( + grad_output.cpu(), + input.cpu(), + weight.cpu(), + bias_sizes, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + output_mask); + + auto to_dev = [&](const at::Tensor& t) { + return t.defined() ? t.to(input.device()) : t; + }; + return ::std::make_tuple( + to_dev(::std::get<0>(out)), + to_dev(::std::get<1>(out)), + to_dev(::std::get<2>(out))); +} +#endif // USE_BPU + // ============================================================ // Generated wrappers for 71 CUDA operators // ============================================================ @@ -347,6 +427,25 @@ TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) { #if defined(FLAGOS_MUSA_KERNEL) #include "backends/musa/generated/musa_register.inc" #endif + #elif defined(USE_BPU) + // BPU registers no compute ops. The BPU's unit of execution is a whole + // compiled graph (a .hbm produced by hbdk4), so there is no per-operator + // kernel to claim -- and claiming an op on PrivateUse1 without a kernel + // behind it raises "backend not registered" instead of falling back. + // Leaving the list out routes every op to cpu_fallback below; acceleration + // comes from the torch.compile backend in torch_fl/backends/bpu/. + // + // The *_overrideable ops are the exception, and the generic fallback cannot + // serve them. aten::convolution routes PrivateUse1 to + // convolution_overrideable, whose only other kernel is a + // CompositeExplicitAutograd stub that raises NotImplementedError -- so + // cpu_fallback, which moves the arguments to CPU and redispatches the *same* + // op, lands right back on that stub. These wrappers instead call + // at::convolution on the CPU tensors, which is the op that actually has a + // CPU kernel. + m.impl("convolution_overrideable", BPUWrapperConvolutionOverrideable); + m.impl("convolution_backward_overrideable", + BPUWrapperConvolutionBackwardOverrideable); #else #define FLAGOS_GEN_IMPLS #include "generated/register.inc" diff --git a/csrc/runtime/accelerator/CMakeLists.txt b/csrc/runtime/accelerator/CMakeLists.txt index 998bea55..d523f2fc 100644 --- a/csrc/runtime/accelerator/CMakeLists.txt +++ b/csrc/runtime/accelerator/CMakeLists.txt @@ -20,6 +20,8 @@ elseif(ACCELERATOR STREQUAL "metax") project(FLAGOS_RUNTIME CXX C) elseif(ACCELERATOR STREQUAL "tsingmicro") project(FLAGOS_RUNTIME CXX C) +elseif(ACCELERATOR STREQUAL "bpu") + project(FLAGOS_RUNTIME CXX C) elseif(ACCELERATOR STREQUAL "dcu") project(FLAGOS_RUNTIME CXX C) elseif(ACCELERATOR STREQUAL "gcu") @@ -112,6 +114,31 @@ elseif(ACCELERATOR STREQUAL "musa") ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_SOURCE_DIR} ${MUSA_HOME}/include) target_link_directories(${LIBRARY_NAME} PRIVATE ${MUSA_HOME}/lib) target_link_libraries(${LIBRARY_NAME} PRIVATE musart) +elseif(ACCELERATOR STREQUAL "bpu") + # D-Robotics RDK BPU. The UCP allocator (libhbucp) provides device memory and + # libbpu the core/task API; both ship in the board image under /usr/hobot/lib + # with headers at /usr/include/hobot, so there is no SDK root to configure. + file(GLOB SOURCE_FILES "${CMAKE_CURRENT_SOURCE_DIR}/bpu/*.cc") + + add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES}) + + find_library(HBUCP_LIB hbucp PATHS ${HOBOT_LIB_DIR} /usr/hobot/lib) + find_library(LIBBPU_LIB bpu PATHS ${HOBOT_LIB_DIR} /usr/hobot/lib) + find_path(HOBOT_UCP_INCLUDE_DIR hb_ucp_sys.h + PATHS ${HOBOT_INCLUDE_DIR} /usr/include/hobot) + + if(NOT HBUCP_LIB OR NOT LIBBPU_LIB OR NOT HOBOT_UCP_INCLUDE_DIR) + message(FATAL_ERROR + "ACCELERATOR=bpu selected, but the Horizon runtime was not found. " + "Expected libhbucp/libbpu under /usr/hobot/lib and hb_ucp_sys.h under " + "/usr/include/hobot (set HOBOT_LIB_DIR / HOBOT_INCLUDE_DIR to override). " + "These ship in the RDK BPU board image.") + endif() + + target_include_directories(${LIBRARY_NAME} PUBLIC + ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_SOURCE_DIR} + ${HOBOT_UCP_INCLUDE_DIR} ${HOBOT_UCP_INCLUDE_DIR}/dnn) + target_link_libraries(${LIBRARY_NAME} PRIVATE ${HBUCP_LIB} ${LIBBPU_LIB}) else() file(GLOB SOURCE_FILES "${CMAKE_CURRENT_SOURCE_DIR}/cuda/*.cc") diff --git a/csrc/runtime/accelerator/bpu/bpu.h b/csrc/runtime/accelerator/bpu/bpu.h new file mode 100644 index 00000000..b7210187 --- /dev/null +++ b/csrc/runtime/accelerator/bpu/bpu.h @@ -0,0 +1,38 @@ +// 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. + +// BPU-specific extensions beyond the flagos runtime contract. + +#pragma once + +#include + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +// Physical address backing a device pointer returned by Malloc, or 0 if the +// pointer is not BPU device memory. +// +// The BPU addresses tensors by physical address (hbUCPSysMem.phyAddr) while +// torch only tracks the host virtual pointer, so bridging the two is what lets +// the compiled .hbm read a torch tensor's storage in place instead of going +// through a numpy copy at each partition boundary. +FLAGOS_EXPORT uint64_t FlagosBPUPhysicalAddress(const void* ptr); + +#ifdef __cplusplus +} // extern "C" +#endif diff --git a/csrc/runtime/accelerator/bpu/device.cc b/csrc/runtime/accelerator/bpu/device.cc new file mode 100644 index 00000000..5a774133 --- /dev/null +++ b/csrc/runtime/accelerator/bpu/device.cc @@ -0,0 +1,84 @@ +// 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. + +// BPU device layer. +// +// The board has four BPU cores (/dev/bpu_core0-3) but they are not independent +// devices: hb_bpu_core_open() takes a core *mask* and a scheduling policy, so +// core selection is the driver's job, not the caller's. Presenting one device +// keeps torch's device indices honest -- claiming four would imply four separate +// memory spaces, which is wrong, since UCP memory is allocated for the SoC. + +#include + +extern "C" { +#include +} + +namespace { + +// Kept for API symmetry with the other backends; only index 0 is ever valid. +thread_local int gCurrentDevice = 0; + +} // namespace + +Error_t GetDeviceCount(int* count) { + if (!count) { + return ErrorUnknown; + } + // Report a device only when the BPU is actually usable. hb_bpu_core_num() + // returns 0 when the driver is absent, which is what torch should see rather + // than a device that fails on first allocation. + *count = (hb_bpu_core_num() > 0) ? 1 : 0; + return Success; +} + +Error_t GetDevice(int* device) { + if (!device) { + return ErrorUnknown; + } + *device = gCurrentDevice; + return Success; +} + +Error_t SetDevice(int device) { + int count = 0; + GetDeviceCount(&count); + if (device < 0 || device >= count) { + return ErrorInvalidDevice; + } + gCurrentDevice = device; + return Success; +} + +Error_t DeviceGetStreamPriorityRange( + int* leastPriority, + int* greatestPriority) { + // The BPU driver does have task priorities (hb_bpu_task_set_prio), but they + // apply per submitted task, not per stream. Report no range. + if (leastPriority) { + *leastPriority = 0; + } + if (greatestPriority) { + *greatestPriority = 0; + } + return Success; +} + +Error_t DeviceSynchronize(void) { + // Task submission goes through hbUCPSubmitTask/hbUCPWaitTaskDone, which the + // runtime already waits on before returning, so there is never outstanding + // work to drain at this level. + return Success; +} diff --git a/csrc/runtime/accelerator/bpu/memory.cc b/csrc/runtime/accelerator/bpu/memory.cc new file mode 100644 index 00000000..37167f93 --- /dev/null +++ b/csrc/runtime/accelerator/bpu/memory.cc @@ -0,0 +1,269 @@ +// 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. + +// D-Robotics RDK BPU (Horizon "nash-p" BPU) memory layer. +// +// The UCP allocator returns memory that is mapped into the calling process as +// well as addressable by the BPU: hbUCPSysMem carries both a host `virAddr` and +// a `phyAddr`. That is what makes a real device allocator possible here -- +// device pointers are host-dereferenceable, so Memcpy is a plain memcpy and no +// bounce buffer is needed. +// +// `MallocCached` is used rather than `hbUCPMalloc` because the host writes +// tensor data through the mapping; the cache has to be flushed before the BPU +// reads it (see the Memcpy paths below), which is cheaper than uncached host +// stores. + +#include + +extern "C" { +#include +} + +#include +#include +#include +#include +#include + +namespace { + +struct Block { + MemoryType type = MemoryType::MemoryTypeUnmanaged; + int device = -1; + hbUCPSysMem mem{}; // phyAddr/virAddr/memSize as returned by the UCP allocator + size_t size = 0; +}; + +class MemoryManager { + public: + static MemoryManager& getInstance() { + static MemoryManager instance; + return instance; + } + + Error_t allocate(void** ptr, size_t size, MemoryType type) { + if (!ptr || size == 0) { + return ErrorUnknown; + } + + int device = 0; + if (type == MemoryType::MemoryTypeDevice) { + GetDevice(&device); + } + + hbUCPSysMem mem{}; + // Host pinned memory has no separate UCP API; the device mapping is already + // host-visible, so both kinds come from the same allocator. The distinction + // is kept in the registry so PointerGetAttributes stays truthful. + int32_t rc = hbUCPMallocCached(&mem, static_cast(size), device); + if (rc != 0 || mem.virAddr == nullptr) { + fprintf( + stderr, + "[flagos-bpu] hbUCPMallocCached(%zu bytes, device %d) failed: rc=%d\n", + size, + device, + static_cast(rc)); + *ptr = nullptr; + return ErrorMemoryAllocation; + } + + { + std::lock_guard lock(m_mutex); + m_registry[mem.virAddr] = Block{type, device, mem, size}; + } + *ptr = mem.virAddr; + return Success; + } + + Error_t free(void* ptr) { + if (!ptr) { + return Success; + } + + hbUCPSysMem mem{}; + { + std::lock_guard lock(m_mutex); + auto it = m_registry.find(ptr); + if (it == m_registry.end()) { + return ErrorUnknown; + } + mem = it->second.mem; + m_registry.erase(it); + } + + return (hbUCPFree(&mem) == 0) ? Success : ErrorUnknown; + } + + // Flush the CPU cache for whichever managed block contains [ptr, ptr+count). + // Cache maintenance is per-allocation in the UCP API, so a partial write + // still flushes its whole block; that is correct, just conservative. + Error_t flush(const void* ptr, int flag) { + std::lock_guard lock(m_mutex); + Block* info = getBlockInfoNoLock(ptr); + if (!info) { + // Not ours (e.g. a plain host buffer) -- nothing to maintain. + return Success; + } + return (hbUCPMemFlush(&info->mem, flag) == 0) ? Success : ErrorUnknown; + } + + Error_t memcpy_(void* dst, const void* src, size_t count, MemcpyKind kind) { + if (!dst || !src || count == 0) { + return ErrorUnknown; + } + + // Device memory is host-mapped, so every direction is a host memcpy. What + // differs is the cache maintenance around it: invalidate before reading + // what the BPU wrote, clean after writing what the BPU will read. + switch (kind) { + case MemcpyDeviceToHost: + case MemcpyDeviceToDevice: + flush(src, HB_SYS_MEM_CACHE_INVALIDATE); + break; + default: + break; + } + + std::memcpy(dst, src, count); + + switch (kind) { + case MemcpyHostToDevice: + case MemcpyDeviceToDevice: + flush(dst, HB_SYS_MEM_CACHE_CLEAN); + break; + default: + break; + } + return Success; + } + + Error_t getPointerAttributes(PointerAttributes* attributes, const void* ptr) { + if (!attributes || !ptr) { + return ErrorUnknown; + } + + std::lock_guard lock(m_mutex); + Block* info = getBlockInfoNoLock(ptr); + if (!info) { + attributes->type = MemoryType::MemoryTypeUnmanaged; + attributes->device = -1; + attributes->pointer = const_cast(ptr); + } else { + attributes->type = info->type; + attributes->device = info->device; + attributes->pointer = info->mem.virAddr; + } + return Success; + } + + Error_t memset_(void* devPtr, int value, size_t count) { + if (!devPtr || count == 0) { + return ErrorUnknown; + } + std::memset(devPtr, value, count); + flush(devPtr, HB_SYS_MEM_CACHE_CLEAN); + return Success; + } + + // Physical address of a managed pointer, for building hbDNNTensor without a + // copy. Returns 0 when the pointer is not device memory we allocated. + uint64_t physicalAddress(const void* ptr) { + std::lock_guard lock(m_mutex); + Block* info = getBlockInfoNoLock(ptr); + if (!info) { + return 0; + } + const auto offset = static_cast(ptr) - + static_cast(info->mem.virAddr); + return info->mem.phyAddr + static_cast(offset); + } + + private: + MemoryManager() = default; + + Block* getBlockInfoNoLock(const void* ptr) { + auto it = m_registry.upper_bound(const_cast(ptr)); + if (it != m_registry.begin()) { + --it; + const char* p = static_cast(ptr); + const char* base = static_cast(it->first); + if (p >= base && p < (base + it->second.size)) { + return &it->second; + } + } + return nullptr; + } + + std::map m_registry; + std::mutex m_mutex; +}; + +} // namespace + +Error_t Malloc(void** devPtr, size_t size) { + return MemoryManager::getInstance().allocate( + devPtr, size, MemoryType::MemoryTypeDevice); +} + +Error_t Free(void* devPtr) { + return MemoryManager::getInstance().free(devPtr); +} + +Error_t MallocHost(void** hostPtr, size_t size) { + return MemoryManager::getInstance().allocate( + hostPtr, size, MemoryType::MemoryTypeHost); +} + +Error_t FreeHost(void* hostPtr) { + return MemoryManager::getInstance().free(hostPtr); +} + +Error_t Memcpy(void* dst, const void* src, size_t count, MemcpyKind kind) { + return MemoryManager::getInstance().memcpy_(dst, src, count, kind); +} + +Error_t MemcpyAsync( + void* dst, + const void* src, + size_t count, + MemcpyKind kind, + Stream_t /*stream*/) { + // The BPU submission path is synchronous (hbUCPSubmitTask + WaitTaskDone), so + // there is no asynchronous copy engine to target. A synchronous copy is a + // valid implementation of the async contract. + return MemoryManager::getInstance().memcpy_(dst, src, count, kind); +} + +Error_t PointerGetAttributes(PointerAttributes* attributes, const void* ptr) { + return MemoryManager::getInstance().getPointerAttributes(attributes, ptr); +} + +Error_t Memset(void* devPtr, int value, size_t count) { + return MemoryManager::getInstance().memset_(devPtr, value, count); +} + +Error_t MemsetAsync( + void* devPtr, + int value, + size_t count, + Stream_t /*stream*/) { + return MemoryManager::getInstance().memset_(devPtr, value, count); +} + +// Non-contract helper used by the Python runtime to hand a tensor's storage +// straight to hbDNNInferV2. Declared in bpu.h. +extern "C" uint64_t FlagosBPUPhysicalAddress(const void* ptr) { + return MemoryManager::getInstance().physicalAddress(ptr); +} diff --git a/csrc/runtime/accelerator/bpu/stream.cc b/csrc/runtime/accelerator/bpu/stream.cc new file mode 100644 index 00000000..067f6878 --- /dev/null +++ b/csrc/runtime/accelerator/bpu/stream.cc @@ -0,0 +1,145 @@ +// 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. + +// BPU stream and event layer. +// +// The BPU has no stream abstraction. Work is submitted as a task +// (hbUCPSubmitTask) and waited on (hbUCPWaitTaskDone) by the runtime before it +// returns, so every operation is already complete by the time torch could +// observe it. Streams therefore reduce to opaque non-null tokens and events to +// host-side timestamps -- which keeps torch's stream/event machinery working +// (c10::Stream requires a valid handle) without pretending to an ordering +// guarantee the hardware does not expose. +// +// EventElapsedTime is genuinely useful despite this: since submission is +// synchronous, wall-clock between two records is the real device time. + +#include + +#include +#include + +namespace { + +using Clock = std::chrono::steady_clock; + +struct EventImpl { + Clock::time_point stamp{}; + bool recorded = false; +}; + +// A single sentinel token stands in for the (only, implicit) stream. Handing out +// a non-null pointer matters: c10 treats a null stream handle as invalid. +// (not constexpr: a reinterpret_cast is not a constant expression) +const auto kDefaultStream = reinterpret_cast(static_cast(0x1)); + +} // namespace + +Error_t StreamCreateWithPriority( + Stream_t* stream, + unsigned int /*flags*/, + int /*priority*/) { + if (!stream) { + return ErrorUnknown; + } + *stream = kDefaultStream; + return Success; +} + +Error_t StreamCreate(Stream_t* stream) { + return StreamCreateWithPriority(stream, 0, 0); +} + +Error_t StreamGetPriority(Stream_t /*stream*/, int* priority) { + if (!priority) { + return ErrorUnknown; + } + *priority = 0; + return Success; +} + +Error_t StreamDestroy(Stream_t /*stream*/) { + // Nothing was allocated in StreamCreate, so there is nothing to release. + return Success; +} + +Error_t StreamQuery(Stream_t /*stream*/) { + // Submission is synchronous: a stream is never busy when observed. + return Success; +} + +Error_t StreamSynchronize(Stream_t /*stream*/) { + return Success; +} + +Error_t StreamWaitEvent( + Stream_t /*stream*/, + Event_t /*event*/, + unsigned int /*flags*/) { + // Any recorded event has already completed, so the wait is satisfied. + return Success; +} + +Error_t EventCreateWithFlags(Event_t* event, unsigned int /*flags*/) { + if (!event) { + return ErrorUnknown; + } + *event = reinterpret_cast(new EventImpl()); + return Success; +} + +Error_t EventCreate(Event_t* event) { + return EventCreateWithFlags(event, 0); +} + +Error_t EventDestroy(Event_t event) { + delete reinterpret_cast(event); + return Success; +} + +Error_t EventRecord(Event_t event, Stream_t /*stream*/) { + auto* impl = reinterpret_cast(event); + if (!impl) { + return ErrorUnknown; + } + impl->stamp = Clock::now(); + impl->recorded = true; + return Success; +} + +Error_t EventSynchronize(Event_t /*event*/) { + return Success; +} + +Error_t EventQuery(Event_t event) { + auto* impl = reinterpret_cast(event); + if (!impl) { + return ErrorUnknown; + } + return impl->recorded ? Success : ErrorNotReady; +} + +Error_t EventElapsedTime(float* ms, Event_t start, Event_t end) { + auto* a = reinterpret_cast(start); + auto* b = reinterpret_cast(end); + if (!ms || !a || !b) { + return ErrorUnknown; + } + if (!a->recorded || !b->recorded) { + return ErrorNotReady; + } + const std::chrono::duration delta = b->stamp - a->stamp; + *ms = delta.count(); + return Success; +} diff --git a/csrc/runtime/allocator/backends/bpu_memory.h b/csrc/runtime/allocator/backends/bpu_memory.h new file mode 100644 index 00000000..0a862413 --- /dev/null +++ b/csrc/runtime/allocator/backends/bpu_memory.h @@ -0,0 +1,79 @@ +// Copyright (c) 2026, BAAI. All rights reserved. + +#pragma once + +#include "../device_memory_interface.h" + +#include + +namespace c10::flagos { + +// D-Robotics RDK BPU (Horizon BPU) implementation of DeviceMemoryInterface. +// +// Unlike the other backends this one does not call a vendor runtime directly: +// the flagos contract functions in csrc/runtime/accelerator/bpu/ already wrap +// the UCP allocator and own the virtual/physical address registry that the +// zero-copy inference path reads. Going through them keeps a single source of +// truth for which pointers are device memory -- a second, independent +// registry here would disagree the moment either side changed. +class BPUDeviceMemory final : public DeviceMemoryInterface { + public: + Error_t device_malloc(void** ptr, size_t size) override { + Error_t err = Malloc(ptr, size); + if (err != Success) { + fprintf( + stderr, "[flagos-bpu] device_malloc(%zu bytes) failed\n", size); + *ptr = nullptr; + } + return err; + } + + Error_t device_free(void* ptr) override { + return Free(ptr); + } + + Error_t get_device_index(int* device) override { + return GetDevice(device); + } + + Error_t set_device(int device) override { + return SetDevice(device); + } + + Error_t get_memory_info(size_t* free, size_t* total) override { + // The UCP allocator exposes no capacity query, and BPU memory is carved out + // of system DRAM by the ION/CMA pool rather than being a separate device + // heap. Reporting zeroes is the honest answer; the caching allocator only + // uses this for stats, not for allocation decisions. + if (free) { + *free = 0; + } + if (total) { + *total = 0; + } + return Success; + } + + Error_t event_create(Event_t* event) override { + return EventCreate(event); + } + + Error_t event_destroy(Event_t event) override { + return EventDestroy(event); + } + + Error_t event_record(Event_t event, Stream_t stream) override { + return EventRecord(event, stream); + } + + Error_t event_query(Event_t event) override { + return EventQuery(event); + } + + Error_t memcpy(void* dst, const void* src, size_t count, MemcpyKind kind) + override { + return Memcpy(dst, src, count, kind); + } +}; + +} // namespace c10::flagos diff --git a/csrc/runtime/allocator/caching_device_allocator.cc b/csrc/runtime/allocator/caching_device_allocator.cc index 9a6feaa1..7bdb5df5 100644 --- a/csrc/runtime/allocator/caching_device_allocator.cc +++ b/csrc/runtime/allocator/caching_device_allocator.cc @@ -9,7 +9,7 @@ #include "backends/dcu_memory.h" #endif #if !defined(USE_ASCEND) && !defined(USE_TSINGMICRO) && !defined(USE_DCU) && \ - !defined(USE_GCU) && !defined(USE_MUSA) + !defined(USE_GCU) && !defined(USE_MUSA) && !defined(USE_BPU) #include "backends/cuda_memory.h" #endif #if defined(USE_TSINGMICRO) @@ -21,6 +21,9 @@ #if defined(USE_MUSA) #include "backends/musa_memory.h" #endif +#if defined(USE_BPU) +#include "backends/bpu_memory.h" +#endif #include @@ -51,10 +54,18 @@ CachingDeviceAllocator::CachingDeviceAllocator( } CachingDeviceAllocator::~CachingDeviceAllocator() { -#if !defined(USE_TSINGMICRO) +#if !defined(USE_TSINGMICRO) && !defined(USE_BPU) // Release all cached memory on destruction. // On TsingMicro, skip this — the TX runtime may already be shut down // at process exit, causing segfaults in txFree. + // Same on BPU: this allocator is a function-local static, so its destructor + // runs from __run_exit_handlers, by which point libhbucp's own FINI_ARRAY + // teardown may already have released the heap the UCP blocks live in, and + // hbUCPFree aborts with "double free or corruption (fasttop)". Calling + // torch_fl._C._empty_cache() before exit frees the same blocks cleanly, which + // confirms the free path itself is fine and only the exit ordering is not. + // Leaking at process exit is harmless: the kernel reclaims the ION/UCP + // carveout when the fd closes. for (auto& state_ptr : device_states_) { if (state_ptr) { release_cached_blocks(*state_ptr); @@ -543,6 +554,12 @@ CachingDeviceAllocator* GetCachingAllocator() { // registers, so it is deliberately not delegated to (see musa_memory.h). auto backend = std::make_unique(); alloc = std::make_unique(std::move(backend)); +#elif defined(USE_BPU) + // BPU: memory comes from the UCP allocator via the accelerator/bpu + // contract functions, which also own the virtual->physical map the + // zero-copy inference path needs. + auto backend = std::make_unique(); + alloc = std::make_unique(std::move(backend)); #else // CUDA (and Metax, which uses CUDA-compatible API) auto backend = std::make_unique(); diff --git a/csrc/runtime/guard.h b/csrc/runtime/guard.h index 4ead28d6..255329f5 100644 --- a/csrc/runtime/guard.h +++ b/csrc/runtime/guard.h @@ -22,7 +22,7 @@ // Threads toolkit ships no CUDA runtime at all, so the header itself is absent // -- same exclusion as hooks.h and copy_ops.cc already carry. #if !defined(USE_ASCEND) && !defined(USE_TSINGMICRO) && !defined(USE_DCU) && \ - !defined(USE_GCU) && !defined(USE_MUSA) + !defined(USE_GCU) && !defined(USE_MUSA) && !defined(USE_BPU) #define FLAGOS_GUARD_HAS_CUDA_STREAM 1 #include #else diff --git a/csrc/runtime/hooks.h b/csrc/runtime/hooks.h index 2dd00b67..760987c4 100644 --- a/csrc/runtime/hooks.h +++ b/csrc/runtime/hooks.h @@ -13,7 +13,7 @@ #include #if !defined(USE_ASCEND) && !defined(USE_TSINGMICRO) && !defined(USE_GCU) && \ - !defined(USE_MUSA) + !defined(USE_MUSA) && !defined(USE_BPU) #include #endif #if defined(USE_MUSA) @@ -66,7 +66,8 @@ struct HooksInterface : public at::PrivateUse1HooksInterface { if (merr != musaSuccess) { musaGetLastError(); } -#elif !defined(USE_ASCEND) && !defined(USE_TSINGMICRO) && !defined(USE_GCU) +#elif !defined(USE_ASCEND) && !defined(USE_TSINGMICRO) && !defined(USE_GCU) && \ + !defined(USE_BPU) // Fallback: check if it's CUDA pinned memory // This is needed because when CUDA is present, PyTorch's pinned memory // allocator defaults to CUDA's cudaMallocHost, which won't be in flagos's diff --git a/docs/bpu.md b/docs/bpu.md new file mode 100644 index 00000000..54900f1d --- /dev/null +++ b/docs/bpu.md @@ -0,0 +1,281 @@ +# D-Robotics RDK BPU (Horizon BPU) backend + +The BPU is the one platform in this repo where acceleration does **not** come +from operator kernels. Its BPU executes whole compiled graphs, so `torch_fl` +provides a real device (UCP-backed memory, device/stream layer) plus a +`torch.compile` backend, and every eager op runs on the CPU. + +```bash +pip install torch==2.10.0+cpu --index-url https://download.pytorch.org/whl/cpu +ACCELERATOR=bpu pip install --no-build-isolation -e . +``` + +```python +import torch, torch_fl + +model = MyNet().eval() +compiled = torch.compile(model, backend="bpu") +out = compiled(torch.randn(1, 3, 224, 224)) +``` + +## torch version + +**Pinned to the 2.10 series** (`TORCH_PIN = "torch>=2.10,<2.11"` in `setup.py`, +mirrored in `pyproject.toml`'s build requires). This is the same pin every +platform in this repo carries, and it is enforced, not advisory. + +The pin exists because the checked-in `csrc/aten/generated/*` bindings are +generated against one specific ATen surface. A newer torch drifts from them and +fails as a wall of compile errors at build time rather than a clean resolver +error, so the pin is what turns a confusing build break into an install-time +message. Moving to a newer torch is a deliberate act: re-run +`scripts/codegen_ops.py`, do not hand-edit the generated files. + +The board runs `torch 2.10.0+cpu` on `/home/sunrise/miniconda3/bin/python3.14` +(the cp314 aarch64 wheel exists on PyPI). + +## Why not per-op kernels + +`hbDNNInferV2` takes a model handle and a tensor array — a compiled `.hbm` +artifact, not a Conv2d argument list. There is no entry point that computes one +convolution, so there is nothing for a `PrivateUse1` kernel to call. This +mirrors `tsingmicro`: `csrc/aten/register.cc` skips the generated +`register.inc` under `USE_BPU`, and every aten op reaches `cpu_fallback`. + +Convolution is the one exception, and not by choice. `aten::convolution` +dispatches `PrivateUse1` to `convolution_overrideable`, whose only other kernel +is a `CompositeExplicitAutograd` stub that raises `NotImplementedError`. The +boxed fallback cannot rescue it — moving the arguments to CPU and redispatching +the *same* op lands back on the stub — so `register.cc` registers two small +wrappers that cross to CPU and call `at::convolution` instead. Without them +`conv2d` on a flagos tensor raises rather than falling back. + +## Compile pipeline + +``` +Dynamo -> AOTAutograd -> aten FX graph + -> partition (torch_fl/accelerator/bpu/partition.py) + -> freeze weights (params/buffers become ONNX initializers) + -> ONNX export (compiler.py, decompose.py) + -> int8 Q/DQ insertion (qdq.py, calibrate.py) + -> hbdk4 -> .hbm, cached by graph structure + -> hbm_runtime (runtime.py) +``` + +Two parts of this are load-bearing rather than optimizations: + +**Quantization is a precondition.** hbdk4's `convert(advice=True)` says it +outright: `"lower to cpu. P.S. The type of hbir.conv's fin is f32, which should +be si8, si16 on bpu."` A float artifact compiles fine and then runs conv on the +CPU, so the BPU sits idle. Q/DQ insertion is what moves the MAC work onto the +device — measured 6.1 ms to 0.64 ms on a 2-conv net. Set +`FLAGOS_BPU_QUANTIZE=0` for bit-exact float artifacts with no speedup. + +**Weight freezing is what makes offload a win at all.** AOTAutograd lifts every +parameter and buffer to a graph input, so a 2-conv block crosses the partition +boundary with 13 tensors instead of 1 — copied on every call, and emitted as +ONNX graph *inputs* rather than initializers, which also blocks the int8 fold. +Before freezing, BPU offload measured 3x *slower* than eager (25.7 ms vs +2.8 ms); after, 0.849 ms. + +Partitions that cannot be compiled stay in the graph and run eagerly, so a +missing or failing hbdk4 costs performance and never correctness. Pass +`strict=True` to `bpu_backend` to raise instead. + +## Compiling on the board + +hbdk4 ships **x86_64-only** wheels (`hbdk4_compiler-4.7.5-cp310/cp311-manylinux_2_17_x86_64.whl`); +the aarch64 wheels are runtime-only. So compiling on the board means running an +x86_64 Python under an emulator — but it does work, on the **stock kernel**, with +no VM and no cross-compile host. + +One command sets it up: + +```bash +scripts/setup_bpu_hbdk4.sh --wheels /path/to/oe/wheels +export FLAGOS_BPU_X86_PYTHON=~/hbdk4-x86/python/bin/python3.11 +export FLAGOS_BPU_X86_EMULATOR=~/hbdk4-x86/bin/box64 +``` + +Verify with: + +```bash +python -c "from torch_fl.accelerator.bpu.compiler import find_hbdk; print(find_hbdk())" +# -> x86-emul +``` + +The wheels come from the D-Robotics OE package (`oe-package-*.tgz` on +`ftp://oeftp@sdk.d-robotics.cc/`); the script needs `hbdk4_compiler-*cp311*` and +`hbdk4_march-*`. + +### Why the setup is not just "apt install box64" + +Four things have to line up. Each one fails in a way that looks unrelated to the +real cause, so they are worth naming. + +**1. box64 must be built from source; the packaged one is too old.** Debian and +Ubuntu ship 0.2.6 (Jan 2024), which had its page size fixed at build time and +aborts immediately: + +``` +Error: PageSize configuration is wrong: configured with 4096, but got 65536 +``` + +Current box64 reads the host page size at runtime (`box64_pagesize = +sysconf(_SC_PAGESIZE)`) and maps 4 KB-aligned x86 `PT_LOAD` segments onto 64 KB +pages itself. Verified with **v0.4.5**: the same x86_64 binary that 0.2.6 refuses +runs correctly. **No kernel rebuild is needed** — earlier revisions of this doc +described one, and it is unnecessary. + +qemu-user has no equivalent fix. It still fails with `SIGBUS` on any `.so` with +4 KB-aligned segments, and `pip` segfaults under it. box64 is the working path. + +**2. `libhbtl.so` needs an explicit `RTLD_GLOBAL` preload.** `_hbdk*.so` expects +`hbtl::Storage::createExternal` and `hbtl::getStrides` to be resolvable, but does +not list `libhbtl.so` in its own `DT_NEEDED` — it inherits them transitively +through `libHBDKPythonCAPI.so`, which box64 does not reproduce. Without the +preload: + +``` +Error: Symbol _ZN4hbtl7Storage14createExternalE... not found, + cannot apply R_X86_64_JUMP_SLOT +ImportError: Cannot dlopen(".../_hbdk.cpython-311-x86_64-linux-gnu.so") +``` + +`compiler.py` handles this: `x86_env()` puts the `_mlir_libs` directory on +`BOX64_LD_LIBRARY_PATH` (box64 resolves guest libraries through its own search +path, so `RUNPATH=$ORIGIN` is not enough) and the compile driver does the +`ctypes.CDLL(..., RTLD_GLOBAL)` preload. + +**3. numba must be *absent* and stubbed, not installed.** hbdk4's ONNX entry +point imports numba unconditionally, but real numba imports `llvmlite.binding`, +whose x86 LLVM JIT segfaults under box64. The crash is in llvmlite's JIT setup — +not in numba, and not in hbdk4. + +Stubbing it out is safe because hbdk4 only ever *calls* numba for custom/numba +ops: `compile_numba()` returns the module untouched when `has_numba_op()` is +false, which is always true for a graph exported from ONNX standard operators. +The stubs raise if actually invoked, so a graph that genuinely needs numba fails +loudly rather than miscompiling. + +**4. torch is stubbed for the same reason.** hbdk4's tracing module imports torch +at module scope but only uses it for `isinstance` checks and `torch.jit.trace` in +the custom-op path. A real x86_64 torch would work but costs several hundred MB +in the emulated environment for code that never runs. Note the stub is the +*emulated* interpreter's torch, entirely separate from the board's aarch64 torch +that torch_fl itself runs on. + +The stubs live in `~/hbdk4-x86/stubs` and are appended to `PYTHONPATH`, never +prepended, so a real numba or torch installed in the guest would win. + +### One tolerated failure + +hbdk4's `compile()` ends by loading the artifact back through hbrt4 to validate +it, which claims BPU device memory. Under emulation that can fail: + +``` +hbrt4_py.Hbrt4PyError: ... Cannot malloc bpu memory with length 52496 bytes: + AllocError { len: 135168 } +``` + +This happens *after* a complete `.hbm` has been written. The compile driver +tolerates exactly this case — non-empty output file plus an exception from the +validation step — because the artifact is the deliverable and torch_fl then loads +it with the board's native aarch64 runtime, which is a stricter check. + +### Caching + +`find_hbdk()` tries a native import, then a CLI driver, then the emulator +(`FLAGOS_BPU_X86_EMULATOR` first, then `box64`, `qemu-x86_64-static`, +`qemu-x86_64`), caching the result. When none is reachable the backend logs a +warning and leaves every partition on the CPU. + +Emulated compilation is slow, which is why `compile_partition` caches artifacts +under `~/.cache/torch_fl_bpu` keyed by graph structure, input signature, march, +and the activation-scale set. Float and int8 builds never share a cache entry. + +## Calibration + +Without calibration, activations use `FLAGOS_BPU_ACT_SCALE` (default 0.05, +covering roughly ±6.35 — wide enough for post-BN/ReLU activations). For better +accuracy, measure real ranges: + +```python +from torch_fl.accelerator.bpu.calibrate import calibrate_onnx +scales = calibrate_onnx("partition.onnx", samples=[x1, x2, ...]) +``` + +Scales key on **ONNX** tensor names, which have no stable relation to torch +module names, so calibration runs on the exported graph via onnxruntime rather +than on the eager module. `calibrate_module` collects per-module ranges when the +torch-side view is what you want. + +## Environment variables + +| Variable | Default | Description | +|----------|---------|-------------| +| `FLAGOS_BPU_X86_PYTHON` | unset | x86_64 python with `hbdk4-compiler`, run under box64 | +| `FLAGOS_BPU_X86_EMULATOR` | unset | path to a box64 binary; needed because the distro 0.2.6 is too old and a self-built one is not on `PATH` | +| `FLAGOS_BPU_X86_STUBS` | `/../../stubs` | numba/torch import stubs for the emulated interpreter | +| `FLAGOS_BPU_MARCH` | `nash-p` | BPU micro-architecture (`nash-p`=BPU, `nash-e`=S100, `nash-m`=S100P) | +| `FLAGOS_BPU_QUANTIZE` | `1` | int8 Q/DQ insertion; `0` compiles float (bit-exact, no BPU speedup) | +| `FLAGOS_BPU_ACT_SCALE` | `0.05` | fallback activation scale for uncalibrated tensors | +| `FLAGOS_BPU_CACHE` | `~/.cache/torch_fl_bpu` | `.hbm` artifact cache | + +## Runtime notes + +- **Device memory is host-mapped.** `hbUCPMallocCached` returns an + `hbUCPSysMem` carrying both a host-writable `virAddr` and a `phyAddr`, and it + works unprivileged. So `Memcpy` is a plain `memcpy` plus cache maintenance + (invalidate before reading what the BPU wrote, clean after writing what it + will read) — no bounce buffer, and a zero-copy inference path is possible + because a tensor's storage can *be* the BPU's memory. + `FlagosBPUPhysicalAddress()` exposes the virtual-to-physical mapping for + building an `hbDNNTensor` without a copy. +- **Four BPU cores, one device.** `hb_bpu_core_open()` takes a core *mask* with + a scheduling policy, and UCP memory is SoC-wide, so cores are a scheduling + detail rather than separate devices. `GetDeviceCount` returns 1, and 0 when + the driver is absent (`hb_bpu_core_num() == 0`). +- **Submission is synchronous** (`hbUCPSubmitTask` + `WaitTaskDone`), so there + is no async copy engine: `MemcpyAsync` is a synchronous copy, streams are a + single sentinel handle, and event timestamps are real `steady_clock` readings + — which makes `EventElapsedTime` genuinely meaningful here. +- **Cached blocks leak at process exit, deliberately.** The caching allocator is + a function-local static, so its destructor runs from `__run_exit_handlers`, by + which point `libhbucp`'s own `FINI_ARRAY` teardown may have released the heap + the UCP blocks live in — `hbUCPFree` then aborts with `double free or + corruption (fasttop)`. `caching_device_allocator.cc` skips the release under + `USE_BPU` (as it already does for tsingmicro). Calling + `torch_fl._C._empty_cache()` before exit frees the same blocks cleanly, which + is what confirms only the exit ordering is at fault. The kernel reclaims the + carveout when the fd closes. + +## Measured performance + +A 6-layer conv stack at 224x224, quantized with frozen weights, compiled +**on the board** (torch 2.10, box64 v0.4.5): **3.75 ms on the BPU vs 72.06 ms +eager CPU — 19.2x.** An earlier measurement on a slightly different stack gave +4.00 ms vs 94.35 ms (23.6x) with cosine similarity 0.981; the ratio depends on +how much of the graph is convolution. + +On a small 2-conv net the relative error against eager float is ~3%, which is the +expected cost of int8 activations — set `FLAGOS_BPU_QUANTIZE=0` for a bit-exact +float artifact, at the price of the conv work falling back to the CPU. + +Toy networks are a wash (0.849 ms vs 0.805 ms) — the fixed submission cost +dominates, and there is not enough MAC work to amortize it. The offload is worth +it when the graph is genuinely convolution-heavy. + +## Known limitations + +- Single device only; the four BPU cores are not scheduled independently. +- Synchronous execution; `hbDNNInferAsync` is unused. +- int8 only. hbdk4 also supports si16, which would trade throughput for accuracy. +- The zero-copy path is partial. `runtime.py` wraps a flagos tensor's UCP + storage in a numpy array in place (`_device_view`), so there is no + device-to-host copy on the way in — but quantization changes dtype (float32 in + the graph, int8 in the artifact), and that conversion copies. Only a + dtype-matched input is truly copy-free. Outputs always copy: `hbm_runtime.run` + allocates its own arrays. Driving `hbDNNInferV2` directly with tensors built + from `FlagosBPUPhysicalAddress()` would close the remaining gap. +- Calibration is opt-in and manual. diff --git a/pyproject.toml b/pyproject.toml index 2cd524bd..36b54fb7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,7 +13,10 @@ # limitations under the License. [build-system] -requires = ["setuptools>=45", "wheel", "cmake>=3.18", "torch"] +# torch is pinned to 2.10 to match the checked-in generated ATen bindings; the +# same pin is applied to install_requires via TORCH_PIN in setup.py, where the +# reasoning is spelled out. +requires = ["setuptools>=45", "wheel", "cmake>=3.18", "torch>=2.10,<2.11"] build-backend = "setuptools.build_meta" [project] diff --git a/scripts/setup_bpu_hbdk4.sh b/scripts/setup_bpu_hbdk4.sh new file mode 100755 index 00000000..733bdcff --- /dev/null +++ b/scripts/setup_bpu_hbdk4.sh @@ -0,0 +1,298 @@ +#!/usr/bin/env bash +# Set up on-board hbdk4 compilation for the RDK BPU. +# +# hbdk4 (the BPU graph compiler) ships x86_64-only wheels, so it cannot run +# natively on this aarch64 board. This script builds the pieces that let it run +# here anyway, entirely on the board -- no VM, no x86 host, and no kernel +# rebuild: +# +# 1. box64 from source. The distro package (0.2.6) aborts with "PageSize +# configuration is wrong: configured with 4096, but got 65536" because its +# page size was a build-time constant. Current box64 reads the host page +# size at runtime and maps 4 KB-aligned x86 segments onto 64 KB pages +# itself, so the stock kernel is fine. +# 2. A standalone x86_64 CPython 3.11 (python-build-standalone). +# 3. hbdk4-compiler + hbdk4-march wheels into that interpreter. +# 4. Import-only stubs for numba and torch -- see stubs/ for why. +# +# Usage: +# scripts/setup_bpu_hbdk4.sh [--wheels DIR] [--prefix DIR] +# +# --wheels defaults to ~/x86vm/wheels; it must contain +# hbdk4_compiler-*-cp311-*_x86_64.whl and hbdk4_march-*_x86_64.whl from the +# D-Robotics OE package (oe-package-*.tgz on ftp://oeftp@sdk.d-robotics.cc/). +# +# Afterwards, export the two variables the script prints and BPU offload turns +# on automatically: +# export FLAGOS_BPU_X86_PYTHON=/python/bin/python3.11 +# export FLAGOS_BPU_X86_EMULATOR=/bin/box64 +set -euo pipefail + +PREFIX="${HOME}/hbdk4-x86" +WHEELS="${HOME}/x86vm/wheels" +BOX64_SRC="${HOME}/box64-src" + +while [[ $# -gt 0 ]]; do + case "$1" in + --wheels) WHEELS="$2"; shift 2 ;; + --prefix) PREFIX="$2"; shift 2 ;; + -h|--help) sed -n '2,30p' "$0"; exit 0 ;; + *) echo "unknown argument: $1" >&2; exit 2 ;; + esac +done + +say() { printf '\n=== %s\n' "$*"; } + +[[ "$(uname -m)" == "aarch64" ]] || { echo "this script is for the aarch64 board" >&2; exit 1; } + +say "page size: $(getconf PAGE_SIZE) (any value is fine; box64 handles 64K)" + +# ---------------------------------------------------------------- box64 +if [[ -x "${PREFIX}/bin/box64" ]]; then + say "box64 already built: $("${PREFIX}/bin/box64" --version 2>&1 | head -1)" +else + say "building box64 from source (needs cmake, gcc; takes a few minutes)" + [[ -d "${BOX64_SRC}" ]] || git clone --depth 1 https://github.com/ptitSeb/box64.git "${BOX64_SRC}" + mkdir -p "${BOX64_SRC}/build" + ( + cd "${BOX64_SRC}/build" + cmake .. -DARM_DYNAREC=ON -DCMAKE_BUILD_TYPE=RelWithDebInfo + make -j"$(nproc)" + ) + mkdir -p "${PREFIX}/bin" + cp "${BOX64_SRC}/build/box64" "${PREFIX}/bin/box64" + say "box64: $("${PREFIX}/bin/box64" --version 2>&1 | head -1)" +fi + +BOX64="${PREFIX}/bin/box64" + +# ------------------------------------------------------- x86_64 CPython +PY="${PREFIX}/python/bin/python3.11" +if [[ -x "${PY}" ]]; then + say "x86_64 python already present" +else + say "fetching standalone x86_64 CPython 3.11" + mkdir -p "${PREFIX}" + python3 - "$PREFIX" <<'PYEOF' +import json, sys, urllib.request +prefix = sys.argv[1] +api = "https://api.github.com/repos/astral-sh/python-build-standalone/releases?per_page=10" +for rel in json.load(urllib.request.urlopen(api, timeout=120)): + hit = [a for a in rel["assets"] + if "cpython-3.11" in a["name"] + and "x86_64-unknown-linux-gnu-install_only." in a["name"]] + if hit: + print("downloading", hit[0]["name"]) + urllib.request.urlretrieve(hit[0]["browser_download_url"], f"{prefix}/py311.tar.gz") + break +else: + raise SystemExit("no cpython-3.11 x86_64 install_only asset found") +PYEOF + tar xf "${PREFIX}/py311.tar.gz" -C "${PREFIX}" + rm -f "${PREFIX}/py311.tar.gz" +fi + +# Confirm the emulator actually runs it. This is the step the old box64 fails. +say "checking box64 can run the x86_64 interpreter" +"${BOX64}" "${PY}" -c 'import platform; print("guest machine:", platform.machine())' + +# --------------------------------------------------------------- hbdk4 +say "installing hbdk4 into the x86_64 interpreter" +shopt -s nullglob +COMPILER_WHL=("${WHEELS}"/hbdk4_compiler-*cp311*x86_64.whl) +MARCH_WHL=("${WHEELS}"/hbdk4_march-*x86_64.whl) +shopt -u nullglob +if [[ ${#COMPILER_WHL[@]} -eq 0 ]]; then + echo "no hbdk4_compiler cp311 x86_64 wheel in ${WHEELS}" >&2 + echo "get it from the D-Robotics OE package; see docs/bpu.md" >&2 + exit 1 +fi + +"${BOX64}" "${PY}" -m pip install --no-cache-dir "${COMPILER_WHL[0]}" "${MARCH_WHL[@]}" +# onnx is the input format; sympy is imported by hbdk4's opset13 module. +# numba is deliberately NOT installed: it imports llvmlite.binding, whose x86 +# LLVM JIT segfaults under box64. The stubs below stand in for it. +"${BOX64}" "${PY}" -m pip install --no-cache-dir onnx sympy + +# ---------------------------------------------------------------- stubs +say "installing import-only stubs for numba and torch" +STUBS="${PREFIX}/stubs" +mkdir -p "${STUBS}/numba/core" "${STUBS}/torch" + +cat > "${STUBS}/numba/__init__.py" <<'EOF' +"""Import-safe stand-in for numba, used only by hbdk4 under box64. + +hbdk4's ONNX entry point imports numba unconditionally +(hbdk4/compiler/onnx/__init__.py -> hbdk4.compiler.numba.tools), but only +*calls* it when the graph contains a numba op: compile_numba() returns the +module untouched when has_numba_op() is false, which is always the case for a +graph exported from ONNX standard operators. + +Real numba cannot be used here. It imports llvmlite.binding, which loads +libllvmlite.so and initializes an x86 LLVM JIT; that segfaults under box64. The +crash is in llvmlite's JIT setup, not in numba or hbdk4. + +Anything that would actually run raises, so a graph that genuinely needs numba +fails loudly instead of silently miscompiling. +""" + +__version__ = "0.0.0+hbdk4-stub" + + +def _unavailable(name): + def _raise(*_args, **_kwargs): + raise RuntimeError( + f"numba.{name} is unavailable: llvmlite's x86 JIT segfaults under " + "box64, so numba is stubbed out. A graph containing a custom/numba " + "op must be compiled on an x86_64 host instead." + ) + + return _raise + + +njit = _unavailable("njit") +typeof = _unavailable("typeof") +EOF + +cat > "${STUBS}/numba/core/__init__.py" <<'EOF' +"""Stub for numba.core. See ../__init__.py.""" + +from . import extending, types # noqa: F401 +EOF + +cat > "${STUBS}/numba/core/extending.py" <<'EOF' +"""Stub for numba.core.extending. + +hbdk4 uses one name from here, is_jitted, and only to assert that a custom op's +entry function carries @numba.njit. Nothing ever does under this stub, so +returning False makes hbdk4's own assertion fire with its own message rather +than an AttributeError. +""" + + +def is_jitted(_func) -> bool: + return False + + +def register_jitable(*_args, **_kwargs): + def _identity(fn): + return fn + + return _identity +EOF + +cat > "${STUBS}/numba/core/types.py" <<'EOF' +"""Stub for numba.core.types: any attribute resolves to a placeholder.""" + + +class _Placeholder: + def __init__(self, name: str): + self._name = name + + def __call__(self, *_args, **_kwargs): + return self + + def __getitem__(self, _key): + return self + + def __repr__(self) -> str: + return f"" + + +def __getattr__(name: str) -> _Placeholder: + return _Placeholder(name) +EOF + +cat > "${STUBS}/torch/__init__.py" <<'EOF' +"""Import-safe stand-in for torch, used only by hbdk4 under box64. + +hbdk4's ONNX path reaches hbdk4/compiler/numba/trace.py, which imports torch at +module scope. Everything it uses torch for lives in the custom-op tracing path +(isinstance checks and torch.jit.trace), which an ONNX-exported graph never +enters. Installing a real x86_64 torch would add several hundred MB to the +emulated environment for code that never runs. + +This is the *emulated* interpreter's torch, entirely separate from the board's +own aarch64 torch that torch_fl runs on. + +Names used in isinstance() or typing.Union[...] must be real classes, which is +why the fallbacks below produce types rather than raising placeholders -- +hbdk4's modules evaluate their annotations eagerly at import time. +""" + +__version__ = "0.0.0+hbdk4-stub" + + +class Tensor: + """For isinstance() only; no object in the ONNX path is an instance.""" + + +class Graph: + """Named in hbdk4's torch-jit adaptor annotations.""" + + +def _unavailable(name): + def _raise(*_args, **_kwargs): + raise RuntimeError( + f"torch.{name} is unavailable: this is a stub torch that exists " + "only so hbdk4's tracing module can be imported under box64. A " + "graph needing it must be compiled on an x86_64 host." + ) + + return _raise + + +zeros = _unavailable("zeros") +from_numpy = _unavailable("from_numpy") + + +class _AttrIsAType: + """Base whose unknown attributes resolve to fresh classes.""" + + def __getattr__(self, name: str): + cls = type(name, (), {}) + setattr(self, name, cls) + return cls + + +class _Nn(_AttrIsAType): + class Module: + pass + + class functional: # noqa: N801 - mirrors torch.nn.functional + pass + + +class _Jit(_AttrIsAType): + trace = staticmethod(_unavailable("jit.trace")) + + +nn = _Nn() +jit = _Jit() + + +def __getattr__(name: str): + return type(name, (), {}) +EOF + +# ---------------------------------------------------------------- verify +say "verifying hbdk4 imports on the board" +LIBS="$(echo "${PREFIX}"/python/lib/python3.*/site-packages/hbdk4/compiler/_mlir_libs)" +PYTHONPATH="${STUBS}" BOX64_LD_LIBRARY_PATH="${LIBS}" "${BOX64}" "${PY}" -c " +import ctypes, os +ctypes.CDLL(os.path.join('${LIBS}', 'libhbtl.so'), mode=ctypes.RTLD_GLOBAL) +from hbdk4.compiler import compile, convert +from hbdk4.compiler.onnx import export +import onnx +print('hbdk4 OK (onnx', onnx.__version__ + ')') +" + +cat < bool: return bool(os.environ.get("PPU_SDK") or os.environ.get("PPU_HOME")) +# The checked-in csrc/aten/generated/* bindings are generated against a +# specific ATen surface, so torch is pinned to the 2.10 series rather than left +# open. Newer torch drifts from those bindings, and a mismatch shows up as a +# wall of compile errors at build time rather than a clean resolver failure -- +# the pin is what turns that into an install-time message. Moving to a newer +# torch is a deliberate act: re-run scripts/codegen_ops.py, do not hand-edit +# the generated files. +TORCH_PIN = "torch>=2.10,<2.11" + + def _install_requires(): - reqs = ["torch"] + reqs = [TORCH_PIN] # FlagGems (and its Triton) is the default operator source, so it is a hard # runtime dep everywhere it can actually run. Platforms that ship their own # Triton are the exception: pulling PyPI's NVIDIA-targeted triton wheel would diff --git a/tests/unit/bpu/test_decompose.py b/tests/unit/bpu/test_decompose.py new file mode 100644 index 00000000..66589f88 --- /dev/null +++ b/tests/unit/bpu/test_decompose.py @@ -0,0 +1,119 @@ +# 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. + +"""The batch-norm rewrite must preserve semantics and be ONNX-exportable.""" + +from __future__ import annotations + +import operator +from pathlib import Path + +import pytest +import torch +from torch._dynamo.backends.common import aot_autograd + +from torch_fl.accelerator.bpu.compiler import export_onnx +from torch_fl.accelerator.bpu.decompose import decompose_for_onnx + +BN_NO_TRAINING = torch.ops.aten._native_batch_norm_legit_no_training.default + + +class ConvBN(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, 3, padding=1) + self.bn = torch.nn.BatchNorm2d(8) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(self.bn(self.conv(x))) + + +def _aten_graph(mod: torch.nn.Module, x: torch.Tensor): + """Capture the post-AOTAutograd aten graph, as the backend sees it.""" + captured = {} + + def fw(gm, inputs): + captured["gm"] = gm + return gm.forward + + with torch.no_grad(): + torch.compile(mod, backend=aot_autograd(fw_compiler=fw))(x) + return captured["gm"] + + +def test_rewrites_batch_norm() -> None: + gm = _aten_graph(ConvBN().eval(), torch.randn(1, 3, 8, 8)) + before = [n for n in gm.graph.nodes if n.target is BN_NO_TRAINING] + assert before, "expected the functional batch_norm in the aten graph" + + decompose_for_onnx(gm) + + assert not [n for n in gm.graph.nodes if n.target is BN_NO_TRAINING] + assert [n for n in gm.graph.nodes if n.target is torch.ops.aten.batch_norm.default] + # The getitem that unpacked the tuple should be gone too. + assert not [ + n + for n in gm.graph.nodes + if n.op == "call_function" + and n.target is operator.getitem + and n.args + and getattr(n.args[0], "target", None) is BN_NO_TRAINING + ] + + +def test_rewrite_preserves_numerics() -> None: + """The rewritten graph must still reproduce eager output exactly.""" + mod = ConvBN().eval() + x = torch.randn(2, 3, 8, 8) + + captured = {} + + def fw(gm, inputs): + decompose_for_onnx(gm) + captured["gm"] = gm + return gm.forward + + with torch.no_grad(): + ref = mod(x) + got = torch.compile(mod, backend=aot_autograd(fw_compiler=fw))(x) + + assert "gm" in captured + assert not [n for n in captured["gm"].graph.nodes if n.target is BN_NO_TRAINING] + torch.testing.assert_close(got, ref) + + +def test_exports_to_onnx_after_rewrite(tmp_path: Path) -> None: + """Without the rewrite this raises UnsupportedOperatorError.""" + pytest.importorskip("onnx") + mod = ConvBN().eval() + x = torch.randn(1, 3, 8, 8) + gm = _aten_graph(mod, x) + + inputs = [] + for n in gm.graph.nodes: + if n.op != "placeholder": + continue + v = n.meta["val"] + inputs.append(torch.zeros(tuple(v.shape), dtype=v.dtype)) + + out = tmp_path / "m.onnx" + ins, outs = export_onnx(gm, inputs, out) + assert out.exists() and ins and outs + + import onnx + + proto = onnx.load(str(out)) + onnx.checker.check_model(proto) + # hbdk4's adaptor needs shapes for intermediates, which export_onnx adds. + assert len(proto.graph.value_info) > 0 diff --git a/tests/unit/bpu/test_eager_device.py b/tests/unit/bpu/test_eager_device.py new file mode 100644 index 00000000..69a0fc8e --- /dev/null +++ b/tests/unit/bpu/test_eager_device.py @@ -0,0 +1,137 @@ +# 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. + +"""Eager-mode behaviour of the flagos device on a BPU build. + +The BPU has no per-op kernels, so eager correctness here is entirely a question +of the *runtime* layer: does UCP-backed memory allocate, does data survive the +round trip through the host mapping, and does every compute op reach the CPU +fallback. Convolution gets its own tests because it does not go through the +generic fallback -- see the BPUWrapperConvolution* wrappers in +csrc/aten/register.cc for why. +""" + +from __future__ import annotations + +import copy + +import pytest +import torch + +torch_fl = pytest.importorskip("torch_fl") + +pytestmark = pytest.mark.skipif( + torch_fl._build_accelerator() != "bpu", + reason="requires a build with ACCELERATOR=bpu", +) + +DEV = "flagos" + + +def test_device_is_available(): + assert torch.accelerator.device_count() >= 1 + + +def test_roundtrip_preserves_data(): + x = torch.randn(64, 64) + assert torch.equal(x.to(DEV).cpu(), x) + + +def test_elementwise_matches_cpu(): + a, b = torch.randn(32, 32), torch.randn(32, 32) + got = (a.to(DEV) + b.to(DEV)) * 2 - 1 + torch.testing.assert_close(got.cpu(), (a + b) * 2 - 1) + + +def test_matmul_matches_cpu(): + a, b = torch.randn(48, 32), torch.randn(32, 16) + torch.testing.assert_close( + (a.to(DEV) @ b.to(DEV)).cpu(), a @ b, atol=1e-5, rtol=1e-5 + ) + + +def test_reduction_and_scalar_readback(): + x = torch.randn(100) + # .item() goes through _local_scalar_dense, a distinct path from the fallback. + assert x.to(DEV).sum().item() == pytest.approx(x.sum().item(), abs=1e-4) + + +def test_conv2d_forward_matches_cpu(): + """aten::convolution on PrivateUse1 routes to convolution_overrideable, + whose only other kernel raises. Without an explicit wrapper this is a + NotImplementedError rather than a CPU fallback.""" + conv = torch.nn.Conv2d(3, 8, 3, padding=1).eval() + x = torch.randn(2, 3, 16, 16) + expected = conv(x) + got = copy.deepcopy(conv).to(DEV)(x.to(DEV)) + assert got.device.type == DEV + torch.testing.assert_close(got.cpu(), expected, atol=1e-5, rtol=1e-5) + + +def test_conv2d_backward_matches_cpu(): + conv = torch.nn.Conv2d(3, 8, 3, padding=1) + + x_cpu = torch.randn(2, 3, 16, 16, requires_grad=True) + conv(x_cpu).sum().backward() + + dev_conv = copy.deepcopy(conv) + dev_conv.zero_grad() + dev_conv = dev_conv.to(DEV) + x_dev = x_cpu.detach().clone().to(DEV).requires_grad_(True) + dev_conv(x_dev).sum().backward() + + torch.testing.assert_close(x_dev.grad.cpu(), x_cpu.grad, atol=1e-4, rtol=1e-4) + torch.testing.assert_close( + dev_conv.weight.grad.cpu(), conv.weight.grad, atol=1e-4, rtol=1e-4 + ) + torch.testing.assert_close( + dev_conv.bias.grad.cpu(), conv.bias.grad, atol=1e-4, rtol=1e-4 + ) + + +def test_module_forward_matches_cpu(): + model = torch.nn.Sequential( + torch.nn.Conv2d(3, 8, 3, padding=1), + torch.nn.BatchNorm2d(8), + torch.nn.ReLU(), + torch.nn.AdaptiveAvgPool2d(1), + torch.nn.Flatten(), + torch.nn.Linear(8, 4), + ).eval() + x = torch.randn(1, 3, 16, 16) + with torch.no_grad(): + expected = model(x) + got = copy.deepcopy(model).to(DEV)(x.to(DEV)) + torch.testing.assert_close(got.cpu(), expected, atol=1e-4, rtol=1e-4) + + +def test_compile_backend_is_registered(): + from torch._dynamo import list_backends + + assert "bpu" in list_backends(exclude_tags=()) + + +def test_compiled_model_matches_eager(): + """Correct with or without hbdk4: partitions that cannot be compiled stay + on the CPU, so this passes either way -- only the tolerance differs.""" + model = torch.nn.Sequential( + torch.nn.Conv2d(3, 16, 3, padding=1), + torch.nn.BatchNorm2d(16), + torch.nn.ReLU(), + ).eval() + x = torch.randn(1, 3, 32, 32) + with torch.no_grad(): + expected = model(x) + got = torch.compile(model, backend="bpu")(x) + torch.testing.assert_close(got, expected, atol=2e-2, rtol=2e-2) diff --git a/tests/unit/bpu/test_freeze.py b/tests/unit/bpu/test_freeze.py new file mode 100644 index 00000000..2e9ca034 --- /dev/null +++ b/tests/unit/bpu/test_freeze.py @@ -0,0 +1,182 @@ +# 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. + +"""Tests for baking parameters and buffers into the compiled artifact. + +AOTAutograd lifts every parameter and buffer to a graph input, so a partition's +boundary carries all of them unless they are frozen. That costs a host copy per +tensor per call and, worse, makes the ONNX exporter emit weights as graph +inputs, which stops the QDQ pass from folding them to int8. These tests pin the +structural properties; none of them needs hbdk4 or the device. +""" + +from __future__ import annotations + +import pytest +import torch +from torch._dynamo.backends.common import aot_autograd + +import torch_fl.accelerator.bpu.backend as backend_mod +from torch_fl.accelerator.bpu.partition import ( + extract_subgraph, + partition_graph, + runtime_inputs, +) + + +class ConvBN(torch.nn.Module): + def __init__(self): + super().__init__() + self.c1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.b1 = torch.nn.BatchNorm2d(8) + self.c2 = torch.nn.Conv2d(8, 8, 3, padding=1) + + def forward(self, x): + return torch.relu(self.c2(torch.relu(self.b1(self.c1(x))))) + + +def _capture(model, x, min_nodes=2): + """Run the real backend plumbing and hand back the aten-level artifacts.""" + captured = {} + + def outer(gm, inputs): + names = [n.name for n in gm.graph.nodes if n.op == "placeholder"] + token = backend_mod._OUTER_INPUTS.set((names, list(inputs))) + + def fw(aten_gm, aten_inputs): + captured["gm"] = aten_gm + captured["frozen"] = backend_mod._frozen_weights(aten_gm, aten_inputs) + captured["parts"] = partition_graph(aten_gm, min_nodes=min_nodes) + return aten_gm.forward + + try: + return aot_autograd(fw_compiler=fw)(gm, inputs) + finally: + backend_mod._OUTER_INPUTS.reset(token) + + with torch.no_grad(): + torch.compile(model, backend=outer)(x) + return captured + + +def test_identifies_parameters_and_buffers(): + cap = _capture(ConvBN().eval(), torch.randn(1, 3, 16, 16)) + # 2 convs (w+b) + BN (w, b, running_mean, running_var) = 8 tensors. + assert len(cap["frozen"]) == 8 + for t in cap["frozen"].values(): + assert isinstance(t, torch.Tensor) + assert not t.is_meta + + +def test_freezing_shrinks_the_boundary_to_activations(): + cap = _capture(ConvBN().eval(), torch.randn(1, 3, 16, 16)) + p = cap["parts"][0] + ri = runtime_inputs(p, cap["frozen"]) + assert len(p.inputs) > len(ri) + assert len(ri) == 1 # just the activation + + +def test_subgraph_signature_matches_runtime_inputs(): + """A mismatch here would pass the wrong tensors to the compiled artifact.""" + cap = _capture(ConvBN().eval(), torch.randn(1, 3, 16, 16)) + p, frozen = cap["parts"][0], cap["frozen"] + sub = extract_subgraph(cap["gm"], p, frozen) + + phs = [n.name for n in sub.graph.nodes if n.op == "placeholder"] + assert phs == [n.name for n in runtime_inputs(p, frozen)] + assert sum(1 for n in sub.graph.nodes if n.op == "get_attr") == len(frozen) + + +def test_extract_subgraph_leaves_the_parent_untouched(): + """Constants are staged on the parent module; they must not linger.""" + cap = _capture(ConvBN().eval(), torch.randn(1, 3, 16, 16)) + gm, p, frozen = cap["gm"], cap["parts"][0], cap["frozen"] + before = set(dict(gm.named_buffers())) | set(gm.__dict__) + extract_subgraph(gm, p, frozen) + leaked = [ + k + for k in (set(dict(gm.named_buffers())) | set(gm.__dict__)) - before + if k.startswith("_frozen_") + ] + assert not leaked + + +def test_frozen_subgraph_is_numerically_equivalent(): + """Freezing must not change what the subgraph computes.""" + cap = _capture(ConvBN().eval(), torch.randn(1, 3, 16, 16)) + gm, p, frozen = cap["gm"], cap["parts"][0], cap["frozen"] + + plain = extract_subgraph(gm, p) + baked = extract_subgraph(gm, p, frozen) + + with torch.no_grad(): + args = [ + frozen[n.name] + if n.name in frozen + else torch.randn(tuple(n.meta["val"].shape), dtype=n.meta["val"].dtype) + for n in p.inputs + ] + want = plain(*args) + got = baked(*[a for n, a in zip(p.inputs, args) if n.name not in frozen]) + + for w, g in zip( + want if isinstance(want, (tuple, list)) else (want,), + got if isinstance(got, (tuple, list)) else (got,), + ): + torch.testing.assert_close(g, w) + + +def test_frozen_weights_export_as_onnx_initializers(): + """The point of freezing: weights must be initializers so QDQ can fold them.""" + onnx = pytest.importorskip("onnx") + from pathlib import Path + from tempfile import TemporaryDirectory + + from torch_fl.accelerator.bpu.compiler import export_onnx + + cap = _capture(ConvBN().eval(), torch.randn(1, 3, 16, 16)) + gm, p, frozen = cap["gm"], cap["parts"][0], cap["frozen"] + sub = extract_subgraph(gm, p, frozen) + ri = runtime_inputs(p, frozen) + ex = [ + torch.zeros(tuple(n.meta["val"].shape), dtype=n.meta["val"].dtype) for n in ri + ] + + with TemporaryDirectory() as td: + path = Path(td) / "s.onnx" + export_onnx(sub, ex, path) + model = onnx.load(str(path)) + + assert len(model.graph.input) == len(ri) + int8 = [i for i in model.graph.initializer if i.data_type == onnx.TensorProto.INT8] + # Two convs: an int8 weight plus a zero_point each, at minimum. + assert len(int8) >= 2 + + +def test_missing_outer_context_degrades_to_no_freezing(): + """Without the outer inputs we must return {}, not raise.""" + model = ConvBN().eval() + captured = {} + + def fw(aten_gm, aten_inputs): + captured["frozen"] = backend_mod._frozen_weights(aten_gm, aten_inputs) + return aten_gm.forward + + with torch.no_grad(): + # aot_autograd used directly, so _OUTER_INPUTS is never set. + torch.compile(model, backend=aot_autograd(fw_compiler=fw))( + torch.randn(1, 3, 16, 16) + ) + + assert captured["frozen"] == {} diff --git a/tests/unit/bpu/test_partition.py b/tests/unit/bpu/test_partition.py new file mode 100644 index 00000000..8587ec07 --- /dev/null +++ b/tests/unit/bpu/test_partition.py @@ -0,0 +1,170 @@ +# 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. + +"""Partitioner tests. These run without hbdk4 — they check graph analysis only.""" + +import logging + +import torch + +from torch_fl.accelerator.bpu.partition import ( + extract_subgraph, + partition_graph, + summarize, +) + + +def _aot_graph(model, *inputs): + """Capture the post-dispatch aten graph the backend actually sees. + + Dynamo alone yields torch-level calls; the partitioner runs after + AOTAutograd lowers them to aten, so tests must capture at the same stage. + """ + from torch._dynamo.backends.common import aot_autograd + + captured = {} + + def grab(gm, example_inputs): + captured.setdefault("gm", gm) + captured.setdefault("inputs", example_inputs) + return gm.forward + + with torch.no_grad(): + torch.compile(model, backend=aot_autograd(fw_compiler=grab))(*inputs) + return captured["gm"], captured["inputs"] + + +class ConvNet(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 16, 3, padding=1) + self.bn = torch.nn.BatchNorm2d(16) + self.pool = torch.nn.MaxPool2d(2) + + def forward(self, x): + return self.pool(torch.relu(self.bn(self.conv(x)))) + + +class WithControlFlow(torch.nn.Module): + """Data-dependent branch forces Dynamo to break the graph.""" + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, 3, padding=1) + self.fc = torch.nn.Linear(8, 4) + + def forward(self, x): + x = torch.relu(self.conv(x)) + if x.sum() > 0: # graph break + x = x * 2 + x = x.mean(dim=(2, 3)) + return self.fc(x) + + +def test_conv_net_forms_one_partition(): + """conv+bn+relu must fuse into a single offloadable region.""" + model = ConvNet().eval() + gm, _ = _aot_graph(model, torch.randn(1, 3, 32, 32)) + parts = partition_graph(gm, min_nodes=2) + + assert len(parts) == 1, summarize(gm, parts) + p = parts[0] + targets = {str(n.target) for n in p.nodes} + assert any("convolution" in t for t in targets) + assert any("relu" in t for t in targets) + # Weights and inputs both cross the boundary; one activation comes out. + assert len(p.inputs) >= 1 + assert len(p.outputs) == 1 + + +def test_subgraph_is_extractable_and_numerically_equal(): + model = ConvNet().eval() + x = torch.randn(1, 3, 32, 32) + gm, _ = _aot_graph(model, x) + parts = partition_graph(gm, min_nodes=2) + sub = extract_subgraph(gm, parts[0]) + + # The extracted subgraph must be a valid, runnable module. + ex = [ + torch.zeros(tuple(n.meta["val"].shape), dtype=n.meta["val"].dtype) + for n in parts[0].inputs + ] + with torch.no_grad(): + out = sub(*ex) + assert out is not None + + +def test_min_nodes_filters_small_partitions(): + class Tiny(torch.nn.Module): + def forward(self, x): + return x + 1 + + gm, _ = _aot_graph(Tiny(), torch.randn(4)) + assert partition_graph(gm, min_nodes=3) == [] + + +def test_control_flow_does_not_break_partitioning(): + model = WithControlFlow().eval() + gm, _ = _aot_graph(model, torch.randn(2, 3, 16, 16)) + # Dynamo hands us one subgraph at a time; partitioning must succeed on it. + parts = partition_graph(gm, min_nodes=2) + assert isinstance(parts, list) + for p in parts: + assert p.nodes + assert p.outputs + + +def test_unsupported_op_splits_partitions(): + """An unsupported op in the middle must yield two separate regions.""" + + class Split(torch.nn.Module): + def forward(self, x): + x = torch.relu(x) + x = x * 2 + x = torch.erfinv(x) # not on the whitelist + x = torch.relu(x) + return x * 3 + + gm, _ = _aot_graph(Split(), torch.randn(8, 8)) + parts = partition_graph(gm, min_nodes=2) + assert len(parts) == 2, summarize(gm, parts) + # Nothing unsupported may leak into a partition. + for p in parts: + assert not any("erfinv" in str(n.target) for n in p.nodes) + + +def test_backend_falls_back_without_hbdk(caplog, monkeypatch): + """Without hbdk4 the model must still produce bit-exact results. + + find_hbdk is stubbed out, because on a machine that *can* compile this would + offload to the BPU and the int8 result would only match approximately — a + different property, covered by the accuracy tests. + """ + import torch_fl.accelerator.bpu.backend as backend_mod + from torch_fl.accelerator.bpu.backend import bpu_backend + + monkeypatch.setattr(backend_mod, "find_hbdk", lambda: None) + + model = ConvNet().eval() + x = torch.randn(1, 3, 32, 32) + with torch.no_grad(): + expected = model(x) + + with caplog.at_level(logging.INFO, logger="torch_fl.bpu"): + compiled = torch.compile(model, backend=bpu_backend) + with torch.no_grad(): + got = compiled(x) + + torch.testing.assert_close(got, expected) + assert "no hbdk4 compiler reachable" in caplog.text diff --git a/tests/unit/bpu/test_qdq.py b/tests/unit/bpu/test_qdq.py new file mode 100644 index 00000000..02fb7eba --- /dev/null +++ b/tests/unit/bpu/test_qdq.py @@ -0,0 +1,222 @@ +# 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. + +"""Tests for QDQ insertion and calibration. + +The property that matters is structural: hbdk4 only puts a conv on the BPU when +its input type is si8/si16, and Q/DQ pairs around the conv are how that is +expressed. These tests check the rewrite is well-formed and numerically close, +without needing the compiler or the device. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +onnx = pytest.importorskip("onnx") + +from torch_fl.accelerator.bpu.calibrate import ( # noqa: E402 + Calibration, + TensorRange, + calibrate_module, +) +from torch_fl.accelerator.bpu.qdq import quantize_onnx # noqa: E402 + + +class TwoConv(torch.nn.Module): + def __init__(self): + super().__init__() + self.c1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.c2 = torch.nn.Conv2d(8, 8, 3, padding=1) + + def forward(self, x): + return torch.relu(self.c2(torch.relu(self.c1(x)))) + + +def _export(tmp_path, model, x): + path = tmp_path / "m.onnx" + with torch.no_grad(): + torch.onnx.export( + model, + (x,), + str(path), + input_names=["in_0"], + output_names=["out_0"], + dynamo=False, + opset_version=17, + ) + proto = onnx.shape_inference.infer_shapes( + onnx.load(str(path)), strict_mode=False, data_prop=True + ) + onnx.save(proto, str(path)) + return path, proto + + +def test_inserts_qdq_around_every_conv(tmp_path): + _, proto = _export(tmp_path, TwoConv().eval(), torch.randn(1, 3, 16, 16)) + n_conv = sum(1 for n in proto.graph.node if n.op_type == "Conv") + assert n_conv == 2 + + out = quantize_onnx(proto) + ops = [n.op_type for n in out.graph.node] + + # One Q + one DQ per activation edge, plus one DQ per weight. + assert ops.count("QuantizeLinear") == n_conv + assert ops.count("DequantizeLinear") == 2 * n_conv + assert ops.count("Conv") == n_conv + onnx.checker.check_model(out) + + +def test_conv_consumes_the_dequantized_edge(tmp_path): + """A Q/DQ pair only helps if the conv actually reads its output.""" + _, proto = _export(tmp_path, TwoConv().eval(), torch.randn(1, 3, 16, 16)) + out = quantize_onnx(proto) + + produced_by = {o: n for n in out.graph.node for o in n.output} + convs = [n for n in out.graph.node if n.op_type == "Conv"] + assert convs + for conv in convs: + assert produced_by[conv.input[0]].op_type == "DequantizeLinear" + assert produced_by[conv.input[1]].op_type == "DequantizeLinear" + + +def test_weights_become_int8_with_signed_zero_point(tmp_path): + """hbdk4's frontend rejects unsigned targets, so zero_point must be int8.""" + from onnx import numpy_helper + + _, proto = _export(tmp_path, TwoConv().eval(), torch.randn(1, 3, 16, 16)) + out = quantize_onnx(proto) + inits = {i.name: numpy_helper.to_array(i) for i in out.graph.initializer} + + q_inputs = [ + n.input + for n in out.graph.node + if n.op_type in ("QuantizeLinear", "DequantizeLinear") + ] + assert q_inputs + for args in q_inputs: + # Scale and zero_point must be constant initializers, not computed. + assert args[1] in inits + assert args[2] in inits + assert inits[args[1]].dtype == np.float32 + assert inits[args[2]].dtype == np.int8 + + int8_weights = [v for v in inits.values() if v.dtype == np.int8 and v.ndim == 4] + assert len(int8_weights) == 2 + for w in int8_weights: + assert w.min() >= -127 and w.max() <= 127 + + +def test_quantized_graph_stays_numerically_close(tmp_path): + ort = pytest.importorskip("onnxruntime") + + model = TwoConv().eval() + x = torch.randn(1, 3, 16, 16) + path, proto = _export(tmp_path, model, x) + + ref = ort.InferenceSession( + onnx.load(str(path)).SerializeToString(), providers=["CPUExecutionProvider"] + ).run(None, {"in_0": x.numpy()})[0] + + out = quantize_onnx(proto, default_act_scale=float(x.abs().max()) / 127.0) + got = ort.InferenceSession( + out.SerializeToString(), providers=["CPUExecutionProvider"] + ).run(None, {"in_0": x.numpy()})[0] + + a, b = ref.ravel(), got.ravel() + cos = float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-12)) + assert cos > 0.99, f"cosine similarity {cos}" + + +def test_no_quantizable_ops_leaves_graph_unchanged(tmp_path): + class OnlyRelu(torch.nn.Module): + def forward(self, x): + return torch.relu(x) * 2 + + _, proto = _export(tmp_path, OnlyRelu().eval(), torch.randn(4, 4)) + before = [n.op_type for n in proto.graph.node] + out = quantize_onnx(proto) + assert [n.op_type for n in out.graph.node] == before + + +def test_calibration_scale_covers_observed_range(): + r = TensorRange() + r.observe(torch.tensor([-2.0, 1.0])) + r.observe(torch.tensor([0.0, 6.5])) # exactly representable in float32 + # Scale must map the widest magnitude to at most QMAX. + assert r.valid + assert 6.5 / r.scale() <= 127.0 + assert r.scale() == pytest.approx(6.5 / 127.0) + + +def test_calibration_handles_degenerate_input(): + empty = TensorRange() + assert not empty.valid + assert empty.scale() == 1.0 + + zeros = TensorRange() + zeros.observe(torch.zeros(4)) + assert zeros.scale() == 1.0 # never 0, which would divide by zero + + +def test_calibrate_module_records_leaf_outputs(): + model = TwoConv().eval() + cal = calibrate_module(model, [torch.randn(1, 3, 16, 16) for _ in range(3)]) + assert isinstance(cal, Calibration) + # Leaf modules plus the synthetic graph input. + assert {"c1", "c2", "__input__"} <= set(cal.ranges) + assert cal.scale_of("__input__") > 0 + + +def test_calibrate_onnx_scales_the_conv_inputs(tmp_path): + pytest.importorskip("onnxruntime") + from torch_fl.accelerator.bpu.calibrate import calibrate_onnx + + model = TwoConv().eval() + path, proto = _export(tmp_path, model, torch.randn(1, 3, 16, 16)) + + scales = calibrate_onnx(path, [torch.randn(1, 3, 16, 16) for _ in range(3)]) + + # One scale per conv activation edge, keyed by ONNX tensor name. + convs = [n for n in proto.graph.node if n.op_type == "Conv"] + assert set(scales) == {c.input[0] for c in convs} + assert all(s > 0 for s in scales.values()) + + +def test_calibrated_scales_beat_a_bad_default(tmp_path): + """Calibration exists to avoid clipping; show that it does.""" + ort = pytest.importorskip("onnxruntime") + from torch_fl.accelerator.bpu.calibrate import calibrate_onnx + + model = TwoConv().eval() + x = torch.randn(1, 3, 16, 16) * 20.0 # far outside the default's range + path, _ = _export(tmp_path, model, x) + + def cos_against_float(**kw): + ref = ort.InferenceSession( + onnx.load(str(path)).SerializeToString(), + providers=["CPUExecutionProvider"], + ).run(None, {"in_0": x.numpy()})[0] + q = quantize_onnx(onnx.load(str(path)), **kw) + got = ort.InferenceSession( + q.SerializeToString(), providers=["CPUExecutionProvider"] + ).run(None, {"in_0": x.numpy()})[0] + a, b = ref.ravel(), got.ravel() + return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-12)) + + default_only = cos_against_float(default_act_scale=0.05) + calibrated = cos_against_float(act_scales=calibrate_onnx(path, [x])) + assert calibrated > default_only diff --git a/tests/unit/bpu/test_splice.py b/tests/unit/bpu/test_splice.py new file mode 100644 index 00000000..1686537a --- /dev/null +++ b/tests/unit/bpu/test_splice.py @@ -0,0 +1,198 @@ +# 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. + +"""Splice tests: verify graph rewriting is correct independently of hbdk4. + +A stub runtime stands in for the BPU and computes the partition on CPU, so any +numerical difference is a graph-rewriting bug rather than a hardware one. +""" + +import torch + +from torch_fl.accelerator.bpu.backend import _BPUCall, _example_inputs_for, _splice +from torch_fl.accelerator.bpu.partition import extract_subgraph, partition_graph + + +class _StubRuntime: + """Runs the extracted subgraph on CPU, mimicking BPURuntime's interface.""" + + def __init__(self, sub): + self.sub = sub + self.calls = 0 + + def __call__(self, *args): + self.calls += 1 + with torch.no_grad(): + out = self.sub(*args) + return list(out) if isinstance(out, (tuple, list)) else [out] + + +def _offload_with_stub(gm, example_inputs, min_nodes=2): + """Replace every partition with a stub-backed call node.""" + parts = partition_graph(gm, min_nodes=min_nodes) + stubs = [] + for i, p in reversed(list(enumerate(parts))): + ex = _example_inputs_for(p, example_inputs, gm) + if ex is None: + continue + sub = extract_subgraph(gm, p) + stub = _StubRuntime(sub) + stubs.append(stub) + _splice(gm, p, _BPUCall(stub, len(p.outputs)), f"_bpu_{i}") + if stubs: + gm.graph.lint() + gm.recompile() + return gm, stubs + + +def _run(model, *inputs, min_nodes=2): + """Compile with stub offload, return (result, stub_call_count).""" + from torch._dynamo.backends.common import aot_autograd + + state = {"stubs": []} + + def compile_aten(aten_gm, aten_inputs): + gm, stubs = _offload_with_stub(aten_gm, aten_inputs, min_nodes) + state["stubs"].extend(stubs) + return gm.forward + + torch._dynamo.reset() + compiled = torch.compile(model, backend=aot_autograd(fw_compiler=compile_aten)) + with torch.no_grad(): + out = compiled(*inputs) + return out, sum(s.calls for s in state["stubs"]) + + +class ConvNet(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 16, 3, padding=1) + self.bn = torch.nn.BatchNorm2d(16) + + def forward(self, x): + return torch.relu(self.bn(self.conv(x))) + + +class TwoRegions(torch.nn.Module): + """Two offloadable regions separated by an unsupported op.""" + + def __init__(self): + super().__init__() + self.c1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.c2 = torch.nn.Conv2d(8, 8, 3, padding=1) + + def forward(self, x): + x = torch.relu(self.c1(x)) * 2 + x = torch.erfinv(x.clamp(-0.9, 0.9)) + return torch.relu(self.c2(x)) * 3 + + +def test_single_partition_numerics(): + model = ConvNet().eval() + x = torch.randn(1, 3, 16, 16) + with torch.no_grad(): + expected = model(x) + + got, n_calls = _run(model, x) + assert n_calls == 1, "partition should have executed exactly once" + torch.testing.assert_close(got, expected) + + +def test_two_partitions_numerics(): + model = TwoRegions().eval() + x = torch.randn(2, 3, 12, 12) + with torch.no_grad(): + expected = model(x) + + got, n_calls = _run(model, x) + assert n_calls == 2, f"expected 2 partition calls, got {n_calls}" + torch.testing.assert_close(got, expected) + + +def test_multi_output_partition(): + """A partition whose results feed two different consumers.""" + + class Branch(torch.nn.Module): + def __init__(self): + super().__init__() + self.c = torch.nn.Conv2d(3, 4, 3, padding=1) + + def forward(self, x): + h = torch.relu(self.c(x)) + a = h * 2 + b = torch.erfinv(h.clamp(-0.9, 0.9)) + return a.sum() + b.sum() + + model = Branch().eval() + x = torch.randn(1, 3, 8, 8) + with torch.no_grad(): + expected = model(x) + got, _ = _run(model, x) + torch.testing.assert_close(got, expected) + + +def test_graph_is_valid_after_splice(): + """The rewritten graph must contain a call_module and no orphan nodes.""" + from torch._dynamo.backends.common import aot_autograd + + seen = {} + + def compile_aten(aten_gm, aten_inputs): + gm, stubs = _offload_with_stub(aten_gm, aten_inputs, 2) + seen["gm"] = gm + seen["n"] = len(stubs) + return gm.forward + + torch._dynamo.reset() + with torch.no_grad(): + torch.compile(ConvNet().eval(), backend=aot_autograd(fw_compiler=compile_aten))( + torch.randn(1, 3, 16, 16) + ) + + gm = seen["gm"] + assert seen["n"] == 1 + call_modules = [n for n in gm.graph.nodes if n.op == "call_module"] + assert len(call_modules) == 1 + # The offloaded aten ops must be gone from the outer graph. + assert not any( + "convolution" in str(n.target) + for n in gm.graph.nodes + if n.op == "call_function" + ) + gm.graph.lint() + + +def test_repeated_calls_reuse_partition(): + """Second invocation must hit the compiled graph, not recompile.""" + model = ConvNet().eval() + x = torch.randn(1, 3, 16, 16) + + from torch._dynamo.backends.common import aot_autograd + + stubs = [] + + def compile_aten(aten_gm, aten_inputs): + gm, s = _offload_with_stub(aten_gm, aten_inputs, 2) + stubs.extend(s) + return gm.forward + + torch._dynamo.reset() + compiled = torch.compile(model, backend=aot_autograd(fw_compiler=compile_aten)) + with torch.no_grad(): + a = compiled(x) + b = compiled(x) + + assert len(stubs) == 1, "backend should compile once" + assert stubs[0].calls == 2, "both runs should reuse the same partition" + torch.testing.assert_close(a, b) diff --git a/tests/unit/bpu/test_x86_env.py b/tests/unit/bpu/test_x86_env.py new file mode 100644 index 00000000..2ce3cb2b --- /dev/null +++ b/tests/unit/bpu/test_x86_env.py @@ -0,0 +1,142 @@ +# 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. + +"""The emulated-hbdk4 environment: stub discovery, library path, preload. + +These are the three things that make hbdk4 importable under box64 (see +docs/bpu.md). They are pure path logic, so they are tested against a fake +directory layout rather than a real x86 install -- which also means these tests +run on any platform, not only on a board that has hbdk4 set up. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from torch_fl.accelerator.bpu import compiler as C + + +@pytest.fixture +def fake_x86(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + """A directory tree shaped like scripts/setup_bpu_hbdk4.sh produces.""" + root = tmp_path / "hbdk4-x86" + py = root / "python" / "bin" / "python3.11" + py.parent.mkdir(parents=True) + py.write_text("#!/bin/false\n") + libs = ( + root + / "python" + / "lib" + / "python3.11" + / "site-packages" + / "hbdk4" + / "compiler" + / "_mlir_libs" + ) + libs.mkdir(parents=True) + (libs / "libhbtl.so").write_bytes(b"") + (root / "stubs" / "numba").mkdir(parents=True) + + monkeypatch.setattr(C, "X86_PYTHON", str(py)) + monkeypatch.setattr(C, "X86_STUBS", "") + return root + + +def test_stub_dir_found_next_to_the_x86_python(fake_x86: Path): + # The default location is derived from the interpreter path so that a + # setup_bpu_hbdk4.sh install needs no extra configuration. + assert C._stub_dir() == str(fake_x86 / "stubs") + + +def test_stub_dir_is_none_without_an_x86_python(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(C, "X86_PYTHON", "") + monkeypatch.setattr(C, "X86_STUBS", "") + assert C._stub_dir() is None + + +def test_explicit_stub_dir_wins( + fake_x86: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + other = tmp_path / "elsewhere" + other.mkdir() + monkeypatch.setattr(C, "X86_STUBS", str(other)) + assert C._stub_dir() == str(other) + + +def test_explicit_stub_dir_that_does_not_exist_is_rejected( + fake_x86: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(C, "X86_STUBS", str(tmp_path / "nope")) + assert C._stub_dir() is None + + +def test_mlir_libs_dir_is_discovered(fake_x86: Path): + found = C._mlir_libs_dir() + assert found is not None + # box64 needs this on its own search path, and libhbtl.so must be here for + # the RTLD_GLOBAL preload to work. + assert Path(found).name == "_mlir_libs" + assert (Path(found) / "libhbtl.so").exists() + + +def test_env_carries_library_path_and_stubs(fake_x86: Path): + env = C.x86_env() + assert env["BOX64_LD_LIBRARY_PATH"] == C._mlir_libs_dir() + assert env["PYTHONPATH"] == str(fake_x86 / "stubs") + + +def test_env_appends_stubs_rather_than_replacing_pythonpath( + fake_x86: Path, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setenv("PYTHONPATH", "/pre/existing") + env = C.x86_env() + parts = env["PYTHONPATH"].split(":") + assert parts[0] == "/pre/existing" + # Appended, never prepended: a real numba or torch in the guest's + # site-packages must win over the stubs. + assert parts[-1] == str(fake_x86 / "stubs") + + +def test_env_preserves_an_existing_box64_library_path( + fake_x86: Path, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setenv("BOX64_LD_LIBRARY_PATH", "/opt/other") + env = C.x86_env() + assert env["BOX64_LD_LIBRARY_PATH"].split(":") == [ + C._mlir_libs_dir(), + "/opt/other", + ] + + +def test_driver_preloads_libhbtl(): + # Without this preload _hbdk.so fails to relocate: it needs hbtl symbols + # that it does not list in DT_NEEDED. See _mlir_libs_dir(). + assert "libhbtl.so" in C._X86_DRIVER + assert "RTLD_GLOBAL" in C._X86_DRIVER + + +def test_driver_tolerates_post_compile_validation_failure(): + # hbdk4's compile() loads the artifact back through hbrt4, which can fail + # with AllocError under emulation after the .hbm is already written. + assert "getsize(out_path) > 0" in C._X86_DRIVER + + +def test_emulator_probe_is_skipped_without_an_x86_python( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(C, "X86_PYTHON", "") + monkeypatch.setattr(C, "_emulator", C._UNSET) + assert C.x86_emulator(refresh=True) is None diff --git a/tests/unit/bpu/test_zero_copy.py b/tests/unit/bpu/test_zero_copy.py new file mode 100644 index 00000000..9878d78b --- /dev/null +++ b/tests/unit/bpu/test_zero_copy.py @@ -0,0 +1,108 @@ +# 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. + +"""The zero-copy numpy view over flagos device storage. + +UCP device memory is mapped into this process, so a tensor on the flagos device +can be handed to hbm_runtime as a numpy array that *is* its storage -- no +device-to-host copy. These tests pin that property down, since it is easy to +regress into a silent `.cpu()` that still passes every correctness test. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +torch_fl = pytest.importorskip("torch_fl") + +from torch_fl.accelerator.bpu.runtime import _as_numpy, _device_view # noqa: E402 + +on_bpu = pytest.mark.skipif( + torch_fl._build_accelerator() != "bpu", + reason="requires a build with ACCELERATOR=bpu", +) + + +@on_bpu +def test_view_shares_storage_with_the_tensor(): + t = torch.arange(12, dtype=torch.float32).reshape(3, 4).to("flagos") + arr = _device_view(t) + assert arr is not None + assert arr.__array_interface__["data"][0] == t.data_ptr() + assert arr.shape == (3, 4) + + +@on_bpu +def test_view_reads_the_tensors_values(): + t = torch.arange(6, dtype=torch.float32).reshape(2, 3) + arr = _device_view(t.to("flagos")) + np.testing.assert_array_equal(arr, t.numpy()) + + +@on_bpu +def test_writes_through_the_view_are_visible_to_torch(): + """The mapping is writable, which is what would let an output land straight + in a torch tensor's storage.""" + t = torch.zeros(4, dtype=torch.float32).to("flagos") + arr = _device_view(t) + arr[:] = [1.0, 2.0, 3.0, 4.0] + torch.testing.assert_close(t.cpu(), torch.tensor([1.0, 2.0, 3.0, 4.0])) + + +@on_bpu +@pytest.mark.parametrize( + "dtype", [torch.float32, torch.int8, torch.uint8, torch.int32, torch.int64] +) +def test_supported_dtypes_get_a_view(dtype): + t = torch.ones(8, dtype=dtype).to("flagos") + arr = _device_view(t) + assert arr is not None + assert arr.__array_interface__["data"][0] == t.data_ptr() + + +@on_bpu +def test_non_contiguous_falls_back(): + """A view would misread a strided tensor, so it must be declined.""" + t = torch.randn(4, 6).to("flagos").t() + assert not t.is_contiguous() + assert _device_view(t) is None + # ...but _as_numpy still produces correct data via the copy path. + np.testing.assert_allclose(_as_numpy(t), t.cpu().numpy()) + + +def test_cpu_tensor_falls_back_to_copy(): + t = torch.randn(3, 3) + assert _device_view(t) is None + arr = _as_numpy(t) + np.testing.assert_allclose(arr, t.numpy()) + + +def test_unsupported_dtype_falls_back(): + t = torch.randn(4, dtype=torch.complex64) + assert _device_view(t) is None + + +@on_bpu +def test_as_numpy_avoids_the_copy_on_device_tensors(): + t = torch.randn(16).to("flagos") + assert _as_numpy(t).__array_interface__["data"][0] == t.data_ptr() + + +@on_bpu +def test_empty_tensor_falls_back(): + """A 0-element buffer has no address worth wrapping.""" + t = torch.empty(0, dtype=torch.float32).to("flagos") + assert _device_view(t) is None diff --git a/torch_fl/__init__.py b/torch_fl/__init__.py index 4863c580..2112da0c 100644 --- a/torch_fl/__init__.py +++ b/torch_fl/__init__.py @@ -1282,6 +1282,41 @@ def _register_compile_backend(): _register_compile_backend() +def _register_bpu_compile_backend() -> None: + """Register torch.compile(backend="bpu") on a BPU build. + + The RDK BPU executes whole compiled graphs (a .hbm produced by hbdk4), not + individual operators, so it has no per-op kernels: eager ops reach + cpu_fallback and all acceleration comes through this backend. That is the + opposite of every other platform here, where the compile path is incidental + and the kernels do the work. + + Import failures are swallowed deliberately. The backend pulls in onnx and + (optionally) hbdk4, so on a board that has the runtime but not the + toolchain, raising here would make `import torch_fl` fail outright and take + the working eager path down with it. + """ + if _build_accelerator() != "bpu": + return + try: + from torch_fl.accelerator import bpu + + bpu.register() + except Exception as exc: # noqa: BLE001 + import warnings + + warnings.warn( + f'torch.compile(backend="bpu") is unavailable: {exc}. ' + "Eager ops still work (they run on the CPU); the BPU offload path " + "needs onnx installed.", + RuntimeWarning, + stacklevel=2, + ) + + +_register_bpu_compile_backend() + + __all__ = [ "flagos", "distributed", diff --git a/torch_fl/accelerator/bpu/__init__.py b/torch_fl/accelerator/bpu/__init__.py new file mode 100644 index 00000000..e1c12f03 --- /dev/null +++ b/torch_fl/accelerator/bpu/__init__.py @@ -0,0 +1,36 @@ +# 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. + +"""D-Robotics RDK BPU support: a torch.compile backend, not per-op kernels. + +The BPU executes whole compiled graphs. An aten subgraph is exported to ONNX, +quantized to int8 (a precondition, not an optimization -- hbdk4 lowers float +conv to the CPU), compiled by hbdk4 into a `.hbm`, and executed through the +Horizon runtime. Ops outside a compilable partition keep running eagerly on the +CPU, so a missing hbdk4 costs performance and never correctness. + +Usage: + + import torch, torch_fl + compiled = torch.compile(model.eval(), backend="bpu") + +`register()` is called by `torch_fl` when the build targets bpu, so importing +this package a second time is not required. +""" + +from __future__ import annotations + +from .backend import bpu_backend, register + +__all__ = ["bpu_backend", "register"] diff --git a/torch_fl/accelerator/bpu/backend.py b/torch_fl/accelerator/bpu/backend.py new file mode 100644 index 00000000..2a5269fb --- /dev/null +++ b/torch_fl/accelerator/bpu/backend.py @@ -0,0 +1,287 @@ +# 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. + +"""The torch.compile backend: partition, compile, and splice BPU calls into the graph. + +Each BPU partition is replaced by one call_function node invoking a BPURuntime, +so the surrounding CPU nodes keep running in eager mode and control flow that +Dynamo already split out is untouched. +""" + +from __future__ import annotations + +import contextvars +import logging +import operator +from typing import Any, Callable + +import torch +from torch.fx import GraphModule, Node + +from .compiler import CompileError, compile_partition, find_hbdk +from .partition import ( + Partition, + extract_subgraph, + partition_graph, + runtime_inputs, + summarize, +) +from .runtime import BPURuntime + +log = logging.getLogger("torch_fl.bpu") + +# Dynamo's own graph inputs, stashed so the inner aten compiler can tell which +# of them are parameters and buffers. By the time AOTAutograd calls us the +# names are gone and the tensors are fake, so the information has to be carried +# across rather than recovered. A ContextVar keeps concurrent compiles apart. +_OUTER_INPUTS: contextvars.ContextVar[tuple[list[str], list[torch.Tensor]] | None] = ( + contextvars.ContextVar("torch_fl_bpu_outer_inputs", default=None) +) + + +class _BPUCall(torch.nn.Module): + """Holds a BPURuntime so it survives as a graph attribute.""" + + def __init__(self, rt: BPURuntime, n_outputs: int): + super().__init__() + self.rt = rt + self.n_outputs = n_outputs + + def forward(self, *args: torch.Tensor): + outs = self.rt(*args) + return outs[0] if self.n_outputs == 1 else tuple(outs) + + +def _example_inputs_for( + p: Partition, + example_inputs: list[torch.Tensor], + gm: GraphModule, + frozen: dict[str, torch.Tensor] | None = None, +) -> list[torch.Tensor] | None: + """Materialize concrete tensors for a partition's boundary inputs. + + Uses the fake tensors Dynamo recorded in node.meta['val'] to synthesize + real tensors of the right shape and dtype. Frozen weights use their real + values, since the compiler bakes those into the artifact; everything else + only needs the right shape. + """ + from torch._subclasses.fake_tensor import unset_fake_temporarily + + frozen = frozen or {} + out = [] + # A backend runs under FakeTensorMode, so torch.zeros would itself produce + # a fake tensor. The ONNX exporter needs real storage. + with unset_fake_temporarily(): + for n in p.inputs: + if n.name in frozen: + continue + val = n.meta.get("val") + if not isinstance(val, torch.Tensor): + return None + if any(not isinstance(d, int) for d in val.shape): + return None + out.append(torch.zeros(tuple(val.shape), dtype=val.dtype)) + return out + + +# Dynamo names lifted module state after its access path, e.g. +# `l_self_modules_c1_parameters_weight_`. AOTAutograd renames these to +# `primals_N` but records the original index in node.meta['desc'].idx, which is +# how a weight is recognised two layers down from where it was named. +_STATE_MARKERS = ("_parameters_", "_buffers_") + + +def _frozen_weights( + gm: GraphModule, example_inputs: list[torch.Tensor] +) -> dict[str, torch.Tensor]: + """Map aten placeholder name -> constant tensor, for parameters and buffers. + + Returns {} when the mapping cannot be established, in which case weights + stay as runtime inputs and the graph is still correct, just slower. + """ + from torch._subclasses.fake_tensor import FakeTensor + + outer = _OUTER_INPUTS.get() + if not outer: + return {} + names, values = outer + frozen: dict[str, torch.Tensor] = {} + + for node in gm.graph.nodes: + if node.op != "placeholder": + continue + idx = getattr(node.meta.get("desc"), "idx", None) + if idx is None or not (0 <= idx < len(values)): + continue + t, name = values[idx], names[idx] + is_state = isinstance(t, torch.nn.Parameter) or any( + m in name for m in _STATE_MARKERS + ) + if is_state and isinstance(t, torch.Tensor) and not isinstance(t, FakeTensor): + frozen[node.name] = t + return frozen + + +def _splice( + gm: GraphModule, + p: Partition, + mod: _BPUCall, + tag: str, + call_inputs: list[Node] | None = None, +) -> None: + """Replace a partition's nodes with a single call into `mod`. + + `call_inputs` must match the compiled artifact's input order — frozen + weights are baked in and are not passed. + """ + setattr(gm, tag, mod) + args = tuple(p.inputs if call_inputs is None else call_inputs) + + with gm.graph.inserting_before(p.nodes[0]): + call = gm.graph.call_module(tag, args=args) + + if len(p.outputs) == 1: + p.outputs[0].replace_all_uses_with(call) + else: + for i, old in enumerate(p.outputs): + with gm.graph.inserting_after(call): + item = gm.graph.call_function(operator.getitem, (call, i)) + old.replace_all_uses_with(item) + + for node in reversed(p.nodes): + if not node.users: + gm.graph.erase_node(node) + + +def bpu_backend( + gm: GraphModule, + example_inputs: list[torch.Tensor], + *, + min_nodes: int = 3, + strict: bool = False, + act_scales: dict[str, float] | None = None, +) -> Callable[..., Any]: + """Dynamo backend that offloads compilable subgraphs to the BPU. + + Dynamo hands us a graph of torch-level calls; AOTAutograd lowers it to + aten ops with fake-tensor metadata, which is what the partitioner matches + against. + + Partitions that cannot be compiled stay in the graph and run on CPU, so a + missing or failing hbdk4 degrades performance but never correctness. Set + `strict=True` to raise instead. + + `act_scales` maps ONNX tensor name to quantization scale (see + `calibrate.calibrate_onnx`). Anything not listed falls back to + `compiler.ACT_SCALE`. + """ + from torch._dynamo.backends.common import aot_autograd + + def _compile_aten(aten_gm: GraphModule, aten_inputs: list[torch.Tensor]): + return _offload( + aten_gm, + aten_inputs, + min_nodes=min_nodes, + strict=strict, + act_scales=act_scales, + ) + + names = [n.name for n in gm.graph.nodes if n.op == "placeholder"] + token = _OUTER_INPUTS.set((names, list(example_inputs))) + try: + return aot_autograd(fw_compiler=_compile_aten)(gm, example_inputs) + finally: + _OUTER_INPUTS.reset(token) + + +def _offload( + gm: GraphModule, + example_inputs: list[torch.Tensor], + *, + min_nodes: int = 3, + strict: bool = False, + act_scales: dict[str, float] | None = None, +) -> Callable[..., Any]: + """Partition an aten graph and splice in BPU calls where possible.""" + partitions = partition_graph(gm, min_nodes=min_nodes) + + if not partitions: + log.info("no BPU-eligible partitions; running on CPU") + return gm.forward + + log.info("%s", summarize(gm, partitions)) + + if find_hbdk() is None: + msg = ( + "no hbdk4 compiler reachable — %d partition(s) will run on CPU. " + "hbdk4 ships x86_64-only wheels, so it runs on this board under " + "box64. Run scripts/setup_bpu_hbdk4.sh once, then set " + "FLAGOS_BPU_X86_PYTHON and FLAGOS_BPU_X86_EMULATOR. The stock " + "64 KB-page kernel is fine (box64 0.4+); see docs/bpu.md." + ) + if strict: + raise CompileError(msg % len(partitions)) + log.warning(msg, len(partitions)) + return gm.forward + + frozen = _frozen_weights(gm, example_inputs) + if frozen: + log.info("freezing %d weight tensor(s) into the artifact", len(frozen)) + else: + log.warning( + "could not identify parameters; weights will cross the boundary on " + "every call, which usually costs more than the offload saves" + ) + + compiled = 0 + # Splice in reverse so earlier partitions' node references stay valid. + for i, p in reversed(list(enumerate(partitions))): + ex = _example_inputs_for(p, example_inputs, gm, frozen) + if ex is None: + log.warning("partition %d: dynamic shapes, keeping on CPU", i) + continue + try: + sub = extract_subgraph(gm, p, frozen) + hbm, ins, outs = compile_partition(sub, ex, act_scales=act_scales) + rt = BPURuntime(str(hbm), ins, outs) + _splice( + gm, + p, + _BPUCall(rt, len(p.outputs)), + f"_bpu_{i}", + runtime_inputs(p, frozen), + ) + compiled += 1 + except CompileError as e: + if strict: + raise + log.warning("partition %d: %s — keeping on CPU", i, e) + except Exception as e: # noqa: BLE001 - never break the user's model + if strict: + raise + log.warning("partition %d: unexpected %s: %s", i, type(e).__name__, e) + + if compiled: + gm.graph.lint() + gm.recompile() + log.info("offloaded %d/%d partition(s) to BPU", compiled, len(partitions)) + + return gm.forward + + +def register() -> None: + """Register the backend so torch.compile(backend="bpu") resolves.""" + from torch._dynamo import register_backend + + register_backend(name="bpu", compiler_fn=bpu_backend) diff --git a/torch_fl/accelerator/bpu/calibrate.py b/torch_fl/accelerator/bpu/calibrate.py new file mode 100644 index 00000000..17aa5e50 --- /dev/null +++ b/torch_fl/accelerator/bpu/calibrate.py @@ -0,0 +1,224 @@ +# 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. + +"""Post-training calibration: collect activation ranges and emit a QDQ ONNX. + +Why this exists +--------------- +The BPU's MAC array is int8/int16. Handing hbdk4 a float graph makes it lower +every conv to `native::Conv2dNHWC` on the CPU — confirmed by `convert(advice=True)`, +which reports verbatim: + + lower to cpu. P.S. The type of hbir.conv's fin is f32, + which should be si8, si16 on bpu. + +So quantization is not an optimization here, it is the precondition for the BPU +participating at all. Measured on a 2-conv net: 6.12 ms (conv on CPU) vs +0.617 ms (conv on BPU), a 10x difference. + +hbdk4's ONNX frontend accepts standard `QuantizeLinear`/`DequantizeLinear` +(opset 13/19) and maps them to its own `qnt.quantize`. Three constraints come +from that frontend and are enforced here: + + * scale must be a constant initializer, not a computed value + * zero_point must be a constant initializer + * the quantized type must be *signed* — int8, never uint8 + +torch 2.13 removed `torch.ao.quantization.quantize_pt2e`, so this uses the +still-present observer machinery directly. That is also easier to control: we +only need per-tensor ranges, not a full quantizer backend config. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, field + +import torch + +log = logging.getLogger("torch_fl.bpu") + +# int8 symmetric range. Using -127 rather than -128 keeps the range symmetric, +# which avoids a bias the BPU's rescale path would otherwise have to absorb. +QMIN, QMAX = -127, 127 + + +@dataclass +class TensorRange: + """Running min/max for one tensor position.""" + + lo: float = float("inf") + hi: float = float("-inf") + n: int = 0 + + def observe(self, t: torch.Tensor) -> None: + if t.numel() == 0: + return + self.lo = min(self.lo, float(t.detach().min())) + self.hi = max(self.hi, float(t.detach().max())) + self.n += 1 + + @property + def valid(self) -> bool: + return self.n > 0 and self.lo <= self.hi + + def scale(self) -> float: + """Symmetric per-tensor scale. + + Symmetric (zero_point=0) is what the BPU's conv path wants; an + asymmetric zero point would need an extra correction term. + """ + if not self.valid: + return 1.0 + m = max(abs(self.lo), abs(self.hi)) + if m == 0.0: + return 1.0 + return m / QMAX + + +@dataclass +class Calibration: + """Activation ranges keyed by tensor name, collected over sample data.""" + + ranges: dict[str, TensorRange] = field(default_factory=dict) + + def observe(self, name: str, t: torch.Tensor) -> None: + self.ranges.setdefault(name, TensorRange()).observe(t) + + def scale_of(self, name: str) -> float | None: + r = self.ranges.get(name) + return r.scale() if r is not None and r.valid else None + + def __len__(self) -> int: + return len(self.ranges) + + def summary(self) -> str: + lines = [f"{len(self.ranges)} tensor(s) calibrated"] + for k, r in list(self.ranges.items())[:8]: + if r.valid: + lines.append(f" {k}: [{r.lo:.4f}, {r.hi:.4f}] scale={r.scale():.6f}") + return "\n".join(lines) + + +def calibrate_onnx( + onnx_path, + samples: list, + quantizable: frozenset[str] = frozenset({"Conv", "Gemm", "MatMul"}), + max_batches: int = 32, +) -> dict[str, float]: + """Return {onnx_tensor_name: scale} for the activations qdq.py will wrap. + + The QDQ pass keys on *ONNX* tensor names, which have no stable relation to + torch module names, so scales have to be measured on the exported graph + rather than on the eager module. Only the first input of each quantizable + op needs a scale; everything else is either a constant weight (scaled from + its own values) or not quantized at all. + + Uses onnxruntime on the float graph. If onnxruntime is missing this returns + {} and the caller falls back to a default scale. + """ + import onnx + + try: + import onnxruntime as ort + except ImportError: + log.warning("onnxruntime unavailable; falling back to a default scale") + return {} + + proto = onnx.load(str(onnx_path)) + g = proto.graph + produced = {o for n in g.node for o in n.output} + initializers = {i.name for i in g.initializer} + + # Activation edges feeding a quantizable op. Graph inputs count; constants + # do not, since dq_weight() derives their scale directly. + wanted = [ + n.input[0] + for n in g.node + if n.op_type in quantizable and n.input and n.input[0] not in initializers + ] + wanted = list(dict.fromkeys(wanted)) + if not wanted: + return {} + + # Intermediates are not graph outputs, so they have to be promoted before + # onnxruntime will hand them back. + existing = {o.name for o in g.output} + for name in wanted: + if name in produced and name not in existing: + g.output.append(onnx.ValueInfoProto(name=name)) + + sess = ort.InferenceSession( + proto.SerializeToString(), providers=["CPUExecutionProvider"] + ) + in_names = [i.name for i in sess.get_inputs()] + out_names = [o.name for o in sess.get_outputs()] + + cal = Calibration() + for sample in samples[:max_batches]: + # A sample is either one tensor (single-input partition) or a sequence + # matching the graph's input order. + row = [sample] if isinstance(sample, torch.Tensor) else list(sample) + feed = { + n: (t.detach().numpy() if isinstance(t, torch.Tensor) else t) + for n, t in zip(in_names, row) + } + outs = sess.run(out_names, feed) + for name, val in zip(out_names, outs): + if name in wanted: + cal.observe(name, torch.from_numpy(val)) + for name, arr in feed.items(): + if name in wanted: + cal.observe(name, torch.as_tensor(arr)) + + scales = {n: s for n in wanted if (s := cal.scale_of(n)) is not None} + log.info("calibration: %d/%d activation scale(s)", len(scales), len(wanted)) + return scales + + +def calibrate_module( + mod: torch.nn.Module, samples: list[torch.Tensor], max_batches: int = 32 +) -> Calibration: + """Run `mod` over `samples`, recording the range of every submodule output. + + Hooks are registered on leaf modules only, so each name corresponds to one + operation rather than a container. + """ + cal = Calibration() + handles = [] + + def hook_for(name: str): + def hook(_m, _inp, out): + if isinstance(out, torch.Tensor): + cal.observe(name, out) + + return hook + + for name, sub in mod.named_modules(): + if name and not list(sub.children()): + handles.append(sub.register_forward_hook(hook_for(name))) + + # The graph input needs a scale too, and no hook fires for it. + mod.eval() + try: + with torch.no_grad(): + for i, x in enumerate(samples[:max_batches]): + cal.observe("__input__", x) + mod(x) + finally: + for h in handles: + h.remove() + + log.info("calibration: %s", cal.summary()) + return cal diff --git a/torch_fl/accelerator/bpu/compiler.py b/torch_fl/accelerator/bpu/compiler.py new file mode 100644 index 00000000..3ebd5227 --- /dev/null +++ b/torch_fl/accelerator/bpu/compiler.py @@ -0,0 +1,495 @@ +# 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. + +"""hbdk4 invocation and compile caching. + +A partition is exported to ONNX, compiled to a .hbm by hbdk4, and the artifact +is cached on disk keyed by graph structure so repeated runs skip compilation. +""" + +from __future__ import annotations + +import hashlib +import logging +import os +import shutil +import subprocess +import tempfile +from pathlib import Path + +import torch +from torch.fx import GraphModule + +log = logging.getLogger("torch_fl.bpu") + +# BPU micro-architecture. The BPU is nash-p; nash-e is S100 and nash-m is +# S100P. Confirmed against the vendor docs and the march string embedded in +# /opt/hobot/model/bpu/*.hbm. +DEFAULT_MARCH = os.environ.get("FLAGOS_BPU_MARCH", "nash-p") + +CACHE_DIR = Path( + os.environ.get("FLAGOS_BPU_CACHE", Path.home() / ".cache" / "torch_fl_bpu") +) + +# Quantization is on by default: without it hbdk4 keeps conv in float and +# lowers it to the CPU, so the BPU never runs the heavy work. Set +# FLAGOS_BPU_QUANTIZE=0 to compile float artifacts (bit-exact, but no speedup). +QUANTIZE = os.environ.get("FLAGOS_BPU_QUANTIZE", "1") not in ("0", "false", "no") + +# Fallback activation scale for tensors with no calibration entry. int8 symmetric +# with this scale covers roughly +-6.35, wide enough for post-BN/ReLU activations. +ACT_SCALE = float(os.environ.get("FLAGOS_BPU_ACT_SCALE", "0.05")) + + +class CompileError(RuntimeError): + """hbdk4 was unavailable or rejected the graph.""" + + +# An x86_64 CPython that has hbdk4 installed, run under an emulator on this +# aarch64 board. hbdk4 ships x86_64-only wheels, so this is the only way to +# compile on the board itself. +X86_PYTHON = os.environ.get("FLAGOS_BPU_X86_PYTHON", "") + +# Explicit emulator override. Useful because the distro box64 is usually too old +# (see x86_emulator) and a self-built one is not on PATH. +X86_EMULATOR = os.environ.get("FLAGOS_BPU_X86_EMULATOR", "") + +# Directory holding import-only stand-ins for numba and torch. hbdk4's ONNX +# entry point imports both unconditionally, but only calls into them for +# custom/numba ops, which a graph exported from ONNX standard operators never +# has. Real numba cannot be used here at all: it imports llvmlite.binding, whose +# x86 LLVM JIT segfaults under box64. Real torch would work but costs several +# hundred MB in the emulated environment for code that never runs. +# +# Defaults to /../stubs so a self-contained setup needs no +# extra configuration; see docs/bpu.md. +X86_STUBS = os.environ.get("FLAGOS_BPU_X86_STUBS", "") + +_UNSET = object() +_emulator: tuple[str, ...] | None | object = _UNSET + + +def _stub_dir() -> str | None: + """Directory of the numba/torch import stubs, or None if absent.""" + if X86_STUBS: + return X86_STUBS if Path(X86_STUBS).is_dir() else None + if not X86_PYTHON: + return None + # ...//python/bin/python3.11 -> ...//stubs + guess = Path(X86_PYTHON).resolve().parent.parent.parent / "stubs" + return str(guess) if guess.is_dir() else None + + +def _mlir_libs_dir() -> str | None: + """hbdk4's bundled shared-library directory inside the x86 environment. + + box64 needs this on BOX64_LD_LIBRARY_PATH, and libhbtl.so needs an explicit + RTLD_GLOBAL preload: _hbdk*.so expects several hbtl symbols + (hbtl::Storage::createExternal, hbtl::getStrides) to be resolvable, but does + not list libhbtl.so in its own DT_NEEDED -- it inherits them transitively + from libHBDKPythonCAPI.so, which box64 does not reproduce. Without the + preload the dlopen fails with "cannot apply R_X86_64_JUMP_SLOT". + """ + if not X86_PYTHON: + return None + root = Path(X86_PYTHON).resolve().parent.parent # .../python + hits = sorted(root.glob("lib/python3.*/site-packages/hbdk4/compiler/_mlir_libs")) + return str(hits[0]) if hits else None + + +def x86_env() -> dict[str, str]: + """Environment for running the x86_64 hbdk4 under an emulator.""" + env = dict(os.environ) + libs = _mlir_libs_dir() + if libs: + # box64 resolves the guest's libraries through its own search path, not + # the host loader's, so RUNPATH=$ORIGIN inside _hbdk.so is not enough. + env["BOX64_LD_LIBRARY_PATH"] = ( + f"{libs}:{env['BOX64_LD_LIBRARY_PATH']}" + if env.get("BOX64_LD_LIBRARY_PATH") + else libs + ) + stubs = _stub_dir() + if stubs: + # Appended, not prepended: a real numba/torch in the guest's + # site-packages should win if one is ever installed there. + env["PYTHONPATH"] = ( + f"{env['PYTHONPATH']}:{stubs}" if env.get("PYTHONPATH") else stubs + ) + return env + + +def x86_emulator(refresh: bool = False) -> tuple[str, ...] | None: + """Command prefix that runs an x86_64 hbdk4 locally, or None. + + Requires FLAGOS_BPU_X86_PYTHON to point at an x86_64 python with + hbdk4-compiler installed, plus an emulator to run it. + + **box64 must be recent.** Its page-size handling used to be a build-time + constant, so the widely packaged 0.2.6 (Debian/Ubuntu) aborts on this board + with "PageSize configuration is wrong: configured with 4096, but got 65536". + Current box64 reads the host page size at runtime (`box64_pagesize = + sysconf(_SC_PAGESIZE)`) and maps 4 KB-aligned x86 PT_LOAD segments onto + 64 KB pages itself, so **the stock 64 KB-page kernel is fine** and no kernel + rebuild is needed. Verified with v0.4.5. Point + FLAGOS_BPU_X86_EMULATOR at a self-built box64 if the packaged one is old. + + qemu-user is still tried, but it has no equivalent fix: it fails with SIGBUS + on any .so whose segments are 4 KB-aligned, and pip segfaults under it. + """ + global _emulator + if _emulator is not _UNSET and not refresh: + return _emulator # type: ignore[return-value] + + _emulator = None + if not X86_PYTHON: + return None + + # Preload libhbtl.so before importing, exactly as the compile driver does, + # so the probe reflects whether a real compile would work. + libs = _mlir_libs_dir() + preload = ( + f"import ctypes; ctypes.CDLL({str(Path(libs) / 'libhbtl.so')!r}, " + "mode=ctypes.RTLD_GLOBAL)\n" + if libs + else "" + ) + probe = preload + "import hbdk4.compiler" + + candidates = [X86_EMULATOR] if X86_EMULATOR else [] + candidates += ["box64", "qemu-x86_64-static", "qemu-x86_64"] + + env = x86_env() + for emul in candidates: + exe = shutil.which(emul) or (emul if Path(emul).is_file() else None) + if not exe: + continue + try: + proc = subprocess.run( + [exe, X86_PYTHON, "-c", probe], + capture_output=True, + text=True, + timeout=900, + env=env, + ) + except (subprocess.TimeoutExpired, OSError): + continue + if proc.returncode == 0: + _emulator = (exe, X86_PYTHON) + log.info("hbdk4 reachable via %s", Path(exe).name) + break + log.debug("%s could not run hbdk4: %s", emul, proc.stderr[-500:]) + + return _emulator # type: ignore[return-value] + + +def find_hbdk() -> str | None: + """Locate a usable hbdk4 compiler driver, if any. + + Returns "python-api" for a native import, "x86-emul" when hbdk4 is only + reachable through an x86_64 emulator (see `x86_emulator`), a path for a CLI + driver, or None. + """ + try: + import hbdk4.compiler # noqa: F401 + + return "python-api" + except ImportError: + pass + + for name in ("hbdk-cc", "hbdk4-cc", "hbdk4-compile", "hb_compile"): + path = shutil.which(name) + if path: + return path + + return "x86-emul" if x86_emulator() else None + + +def graph_key(gm: GraphModule, example_inputs: list[torch.Tensor], march: str) -> str: + """Stable hash over graph structure, input signature and target arch.""" + h = hashlib.sha256() + h.update(march.encode()) + for node in gm.graph.nodes: + h.update(f"{node.op}:{node.target}:".encode()) + val = node.meta.get("val") + if isinstance(val, torch.Tensor): + h.update(f"{tuple(val.shape)}:{val.dtype}".encode()) + for t in example_inputs: + h.update(f"{tuple(t.shape)}:{t.dtype}".encode()) + return h.hexdigest()[:16] + + +def export_onnx( + gm: GraphModule, + example_inputs: list[torch.Tensor], + path: Path, + act_scales: dict[str, float] | None = None, +) -> tuple[list[str], list[str]]: + """Export a partition to ONNX. Returns (input_names, output_names).""" + from .decompose import decompose_for_onnx + + decompose_for_onnx(gm) + + n_out = len(gm.graph.find_nodes(op="output")[0].args[0]) + input_names = [f"bpu_in_{i}" for i in range(len(example_inputs))] + output_names = [f"bpu_out_{i}" for i in range(n_out)] + + gm.eval() + # A Dynamo backend is invoked inside an active FakeTensorMode. The ONNX + # exporter traces the module with real tensors, so that mode has to be + # suspended or every op dispatches to fakes and fails on shape guards. + from torch._subclasses.fake_tensor import unset_fake_temporarily + + with torch.no_grad(), unset_fake_temporarily(): + torch.onnx.export( + gm, + tuple(example_inputs), + str(path), + input_names=input_names, + output_names=output_names, + dynamo=False, + opset_version=17, + ) + + # hbdk4's ONNX adaptor resolves every node output name through value_info, + # so intermediates need inferred shapes or it raises "key ... not found". + try: + import onnx + from onnx import shape_inference + + proto = shape_inference.infer_shapes( + onnx.load(str(path)), strict_mode=False, data_prop=True + ) + + # Without int8 inputs hbdk4 lowers every conv to native::Conv2dNHWC on + # the CPU, so the BPU sits idle. Q/DQ insertion is what moves the MAC + # work onto the device — measured 6.1 ms -> 0.64 ms on a 2-conv net. + if QUANTIZE: + from .qdq import quantize_onnx + + proto = quantize_onnx( + proto, act_scales=act_scales, default_act_scale=ACT_SCALE + ) + + onnx.save(proto, str(path)) + except Exception as e: # noqa: BLE001 + log.warning("onnx post-processing skipped: %s", e) + + return input_names, output_names + + +def compile_hbm(onnx_path: Path, out_path: Path, march: str = DEFAULT_MARCH) -> Path: + """Compile an ONNX file to a .hbm via hbdk4. + + Raises CompileError if the toolchain is missing or the compile fails; the + caller is expected to fall back to CPU execution for that partition. + """ + driver = find_hbdk() + if driver is None: + raise CompileError( + "hbdk4 compiler not found. It ships x86_64-only wheels, so on this " + "aarch64 board point FLAGOS_BPU_X86_PYTHON at an x86_64 python " + "that has hbdk4-compiler installed, and use a recent box64 (0.4.x; " + "the distro 0.2.6 aborts on this board's " + f"{os.sysconf('SC_PAGESIZE')}-byte pages) -- see docs/bpu.md for " + "the one-time setup." + ) + + if driver == "python-api": + return _compile_via_python_api(onnx_path, out_path, march) + + if driver == "x86-emul": + return _compile_via_x86_emulator(onnx_path, out_path, march) + + cmd = [ + driver, + "--model", + str(onnx_path), + "--march", + march, + "-o", + str(out_path), + ] + log.info("hbdk4: %s", " ".join(cmd)) + proc = subprocess.run(cmd, capture_output=True, text=True) + if proc.returncode != 0: + raise CompileError( + f"hbdk4 failed (exit {proc.returncode}):\n" + f"{proc.stderr[-2000:] or proc.stdout[-2000:]}" + ) + if not out_path.exists(): + raise CompileError(f"hbdk4 reported success but {out_path} is missing") + return out_path + + +def _compile_via_python_api( + onnx_path: Path, out_path: Path, march: str, opt: int = 2 +) -> Path: + """Compile through the hbdk4 Python API. + + Signatures per the OE 3.7.0 API reference: + hbdk4.compiler.onnx.export(proto: onnx.ModelProto) -> Module + hbdk4.compiler.convert(m, march) -> Module + hbdk4.compiler.compile(m, path, march, opt=2, ...) -> Hbm + + Note `export` takes a loaded protobuf, not a path, and `compile` writes the + artifact itself — the output path is its second positional argument. + """ + try: + import onnx + from hbdk4.compiler import compile as hbdk_compile + from hbdk4.compiler import convert + from hbdk4.compiler.onnx import export as onnx_export + except ImportError as e: + raise CompileError(f"hbdk4 Python API unavailable: {e}") from e + + try: + proto = onnx.load(str(onnx_path)) + module = onnx_export(proto) + # Lower hbir to this march's backend IR. Any Q/DQ already present in the + # graph (inserted by qdq.py) is what lets conv land on the BPU; without + # it convert() reports "fin is f32, which should be si8, si16 on bpu" + # and falls back to native::Conv2dNHWC on the CPU. + quantized = convert(module, march) + hbdk_compile(quantized, str(out_path), march, opt=opt) + except Exception as e: + raise CompileError(f"hbdk4 compile failed: {type(e).__name__}: {e}") from e + + if not out_path.exists(): + raise CompileError(f"hbdk4 returned but {out_path} is missing") + return out_path + + +# Driver run by the emulated x86 interpreter. It is the same three hbdk4 calls +# as _compile_via_python_api; only the interpreter differs, so keeping it as a +# string avoids shipping a second copy of the logic as a separate file. +# +# The libhbtl.so preload is not optional -- see _mlir_libs_dir() for why. +_X86_DRIVER = """ +import ctypes, os, sys + +libs = os.environ.get("FLAGOS_BPU_MLIR_LIBS") +if libs: + ctypes.CDLL(os.path.join(libs, "libhbtl.so"), mode=ctypes.RTLD_GLOBAL) + +import onnx +from hbdk4.compiler import compile as hbdk_compile, convert +from hbdk4.compiler.onnx import export as onnx_export + +onnx_path, out_path, march, opt = sys.argv[1], sys.argv[2], sys.argv[3], int(sys.argv[4]) +proto = onnx.load(onnx_path) +module = onnx_export(proto) + +try: + hbdk_compile(convert(module, march), out_path, march, opt=opt) +except Exception as e: + # hbdk4's compile() finishes by loading the artifact back through hbrt4 to + # validate it (apis.link -> Hbm(path)), which claims BPU device memory. Under + # emulation that step can fail with ResourceExhausted/AllocError *after* + # link() has already written a complete .hbm. Tolerate exactly that case: the + # file is the deliverable, and the caller loads it with the board's native + # aarch64 runtime anyway, which is a stricter check than this one. + if not (os.path.exists(out_path) and os.path.getsize(out_path) > 0): + raise + print("hbdk4 post-compile validation skipped: %s: %s" % (type(e).__name__, e), + file=sys.stderr) +""" + + +def _compile_via_x86_emulator( + onnx_path: Path, out_path: Path, march: str, opt: int = 2 +) -> Path: + """Compile by running an x86_64 hbdk4 under box64 or qemu-user. + + Everything stays on this machine: the emulator only translates + instructions, so the ONNX input and .hbm output are ordinary local files. + Emulated compilation is slow (no JIT cache across runs), which is why + compile_partition() caches the artifact by graph structure. + """ + emul = x86_emulator() + if emul is None: + raise CompileError("no x86_64 emulator able to run hbdk4") + + cmd = [ + *emul, + "-c", + _X86_DRIVER, + str(onnx_path), + str(out_path), + march, + str(opt), + ] + env = x86_env() + libs = _mlir_libs_dir() + if libs: + # Read by _X86_DRIVER to preload libhbtl.so inside the guest. + env["FLAGOS_BPU_MLIR_LIBS"] = libs + log.info("hbdk4 via %s: %s", Path(emul[0]).name, onnx_path.name) + try: + proc = subprocess.run( + cmd, capture_output=True, text=True, timeout=3600, env=env + ) + except subprocess.TimeoutExpired as e: + raise CompileError("hbdk4 timed out under emulation (1h)") from e + + if proc.returncode != 0: + raise CompileError( + f"hbdk4 failed under emulation (exit {proc.returncode}):\n" + f"{(proc.stderr or proc.stdout)[-2000:]}" + ) + if not out_path.exists(): + raise CompileError(f"hbdk4 returned but {out_path} is missing") + log.info("compiled: %s (%d bytes)", out_path.name, out_path.stat().st_size) + return out_path + + +def compile_partition( + gm: GraphModule, + example_inputs: list[torch.Tensor], + march: str = DEFAULT_MARCH, + cache_dir: Path | None = None, + act_scales: dict[str, float] | None = None, +) -> tuple[Path, list[str], list[str]]: + """Export, compile and cache one partition. + + Returns (hbm_path, input_names, output_names). + """ + cache = cache_dir or CACHE_DIR + cache.mkdir(parents=True, exist_ok=True) + + key = graph_key(gm, example_inputs, march) + # Float and int8 artifacts differ, and so do two int8 builds with different + # activation scales, so none of them may share a cache entry. + if QUANTIZE: + qh = hashlib.sha256(repr(sorted((act_scales or {}).items())).encode()) + qh.update(f"{ACT_SCALE}".encode()) + key = f"{key}q{qh.hexdigest()[:8]}" + hbm_path = cache / f"{key}.hbm" + names_path = cache / f"{key}.names" + + if hbm_path.exists() and names_path.exists(): + ins, outs = names_path.read_text().strip().split("\n") + log.info("cache hit: %s", hbm_path.name) + return hbm_path, ins.split(","), outs.split(",") + + with tempfile.TemporaryDirectory() as td: + onnx_path = Path(td) / f"{key}.onnx" + in_names, out_names = export_onnx( + gm, example_inputs, onnx_path, act_scales=act_scales + ) + compile_hbm(onnx_path, hbm_path, march) + + names_path.write_text(f"{','.join(in_names)}\n{','.join(out_names)}\n") + log.info("compiled: %s", hbm_path.name) + return hbm_path, in_names, out_names diff --git a/torch_fl/accelerator/bpu/decompose.py b/torch_fl/accelerator/bpu/decompose.py new file mode 100644 index 00000000..197a807d --- /dev/null +++ b/torch_fl/accelerator/bpu/decompose.py @@ -0,0 +1,100 @@ +# 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. + +"""Rewrite aten ops that the ONNX exporter cannot handle. + +AOTAutograd emits functional aten variants that the TorchScript-based ONNX +exporter has no symbolic function for. Each rewrite here preserves semantics +exactly and only changes which overload the graph names. +""" + +from __future__ import annotations + +import logging +import operator + +import torch +from torch.fx import GraphModule + +log = logging.getLogger("torch_fl.bpu") + +_BN_NO_TRAINING = torch.ops.aten._native_batch_norm_legit_no_training.default + + +def _rewrite_batch_norm(gm: GraphModule) -> int: + """`_native_batch_norm_legit_no_training` -> `aten.batch_norm`. + + The functional op returns (out, save_mean, save_invstd); in inference only + the first element is ever consumed, and `aten.batch_norm` with + training=False computes exactly that and does have an ONNX symbolic. + """ + changed = 0 + for node in list(gm.graph.nodes): + if node.op != "call_function" or node.target is not _BN_NO_TRAINING: + continue + + users = list(node.users) + # Only the primary output may be used; save_mean/save_invstd are + # training-only statistics and have no ONNX equivalent. + consumed = { + u.args[1] + for u in users + if u.op == "call_function" and u.target is operator.getitem + } + if any( + u.op != "call_function" or u.target is not operator.getitem for u in users + ) or not consumed <= {0}: + log.debug("batch_norm %s: non-trivial users, left alone", node.name) + continue + + inp, weight, bias, running_mean, running_var, momentum, eps = node.args + + with gm.graph.inserting_before(node): + new = gm.graph.call_function( + torch.ops.aten.batch_norm.default, + args=( + inp, + weight, + bias, + running_mean, + running_var, + False, + momentum, + eps, + True, + ), + ) + new.meta.update(node.meta) + val = node.meta.get("val") + if isinstance(val, (tuple, list)) and val: + new.meta["val"] = val[0] + + # Every user is getitem(node, 0); point them straight at `new`. + for u in users: + u.replace_all_uses_with(new) + gm.graph.erase_node(u) + gm.graph.erase_node(node) + changed += 1 + + return changed + + +def decompose_for_onnx(gm: GraphModule) -> GraphModule: + """Apply every rewrite needed before torch.onnx.export. Mutates in place.""" + n = _rewrite_batch_norm(gm) + if n: + log.debug("decompose: rewrote %d batch_norm node(s)", n) + gm.graph.lint() + gm.recompile() + return gm diff --git a/torch_fl/accelerator/bpu/partition.py b/torch_fl/accelerator/bpu/partition.py new file mode 100644 index 00000000..b6dcec27 --- /dev/null +++ b/torch_fl/accelerator/bpu/partition.py @@ -0,0 +1,307 @@ +# 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. + +"""FX graph partitioning: split a graph into BPU-compilable subgraphs and CPU remainder. + +The BPU only executes whole compiled graphs, so the unit of offload is a maximal +connected region of supported nodes rather than an individual operator. +""" + +from __future__ import annotations + +import operator +from dataclasses import dataclass, field +from typing import Callable + +import torch +from torch.fx import Graph, GraphModule, Node + +# Operators hbdk4 can compile. Conservative on purpose: an op that is wrongly +# listed here fails at compile time (recoverable, we fall back), whereas a +# missing op only costs us a smaller subgraph. +_SUPPORTED: set[Callable | str] = { + # convolution / linear + torch.ops.aten.convolution.default, + torch.ops.aten.conv2d.default, + torch.ops.aten.linear.default, + torch.ops.aten.addmm.default, + torch.ops.aten.mm.default, + torch.ops.aten.bmm.default, + # normalization + torch.ops.aten._native_batch_norm_legit_no_training.default, + torch.ops.aten.batch_norm.default, + torch.ops.aten.native_layer_norm.default, + # activation + torch.ops.aten.relu.default, + torch.ops.aten.relu_.default, + torch.ops.aten.hardtanh.default, + torch.ops.aten.hardswish.default, + torch.ops.aten.sigmoid.default, + torch.ops.aten.tanh.default, + torch.ops.aten.silu.default, + torch.ops.aten.gelu.default, + # pooling + torch.ops.aten.max_pool2d.default, + # max_pool2d_with_indices is deliberately excluded: it returns an indices + # tensor the BPU does not produce, and the decomposition is what Dynamo + # actually emits. Use aten.max_pool2d via a graph pass if you need pooling + # inside a partition. + torch.ops.aten.avg_pool2d.default, + torch.ops.aten.mean.dim, + torch.ops.aten.adaptive_avg_pool2d.default, + torch.ops.aten._adaptive_avg_pool2d.default, + # elementwise + torch.ops.aten.add.Tensor, + torch.ops.aten.sub.Tensor, + torch.ops.aten.mul.Tensor, + torch.ops.aten.div.Tensor, + torch.ops.aten.add_.Tensor, + torch.ops.aten.mul_.Tensor, + # shape + torch.ops.aten.view.default, + torch.ops.aten.reshape.default, + torch.ops.aten.permute.default, + torch.ops.aten.transpose.int, + torch.ops.aten.flatten.using_ints, + torch.ops.aten.cat.default, + torch.ops.aten.contiguous.default, + torch.ops.aten.squeeze.dim, + torch.ops.aten.unsqueeze.default, + # misc + torch.ops.aten.softmax.int, + torch.ops.aten._softmax.default, + torch.ops.aten.clone.default, +} + +# Structural nodes that never block a partition: they carry no computation. +_STRUCTURAL = {"placeholder", "output", "get_attr"} + + +@dataclass +class Partition: + """A maximal connected region of BPU-supported nodes.""" + + nodes: list[Node] = field(default_factory=list) + inputs: list[Node] = field(default_factory=list) + outputs: list[Node] = field(default_factory=list) + + def __len__(self) -> int: + return len(self.nodes) + + +def is_supported(node: Node) -> bool: + """Whether a single node can run on the BPU. + + Dynamic shapes are rejected: hbdk4 bakes memspace offsets and SRAM tiling + into the compiled artifact, so every shape must be static at compile time. + """ + if node.op in _STRUCTURAL: + return False + if node.op == "call_method": + return False + if node.op == "call_module": + # Dynamo traces down to aten ops; a surviving call_module is opaque. + return False + if node.op != "call_function": + return False + if node.target in (operator.getitem,): + # Tuple indexing is bookkeeping, absorbed into whichever side needs it. + return True + if node.target not in _SUPPORTED: + return False + + val = node.meta.get("val") + if val is None: + return False + vals = val if isinstance(val, (tuple, list)) else [val] + for v in vals: + if not isinstance(v, torch.Tensor): + continue + if any(not isinstance(d, int) for d in v.shape): + return False # symbolic dim + if v.dtype not in (torch.float32, torch.float16, torch.int8, torch.int32): + return False + return True + + +def partition_graph(gm: GraphModule, min_nodes: int = 3) -> list[Partition]: + """Find maximal connected regions of supported nodes, in topological order. + + A region is grown greedily in graph order: a supported node joins the + current region if all of its supported predecessors are already in it, + which keeps each region a contiguous slice of the topological order and so + directly emittable as a standalone graph. + + Regions smaller than `min_nodes` are dropped — the host/device copy at the + boundary costs more than the offload saves. + """ + supported = {n: is_supported(n) for n in gm.graph.nodes} + + # getitem alone is not worth a partition; it only rides along. + partitions: list[Partition] = [] + current: list[Node] = [] + current_set: set[Node] = set() + + def flush() -> None: + nonlocal current, current_set + if current: + real = [n for n in current if n.target is not operator.getitem] + if real: + partitions.append(Partition(nodes=list(current))) + current = [] + current_set = set() + + for node in gm.graph.nodes: + if not supported.get(node, False): + flush() + continue + # Join only if every supported producer is already inside this region; + # otherwise this node belongs to a later region. + deps_ok = all( + arg in current_set or not supported.get(arg, False) + for arg in node.all_input_nodes + ) + if not deps_ok: + flush() + current.append(node) + current_set.add(node) + flush() + + for p in partitions: + _compute_boundary(p) + + return [p for p in partitions if len(p) >= min_nodes] + + +def _compute_boundary(p: Partition) -> None: + """Fill in the tensors crossing into and out of a partition.""" + inside = set(p.nodes) + + seen_in: set[Node] = set() + for node in p.nodes: + for arg in node.all_input_nodes: + if arg not in inside and arg not in seen_in: + seen_in.add(arg) + p.inputs.append(arg) + + seen_out: set[Node] = set() + for node in p.nodes: + # A node is an output if anything outside the partition consumes it. + if any(user not in inside for user in node.users): + if node not in seen_out: + seen_out.add(node) + p.outputs.append(node) + p.outputs = [n for n in p.outputs if n.target is not operator.getitem] or p.outputs + + +def runtime_inputs( + p: Partition, frozen: dict[str, torch.Tensor] | None = None +) -> list[Node]: + """The partition inputs that must be passed at call time. + + Weights baked into the compiled artifact are excluded — see + `extract_subgraph`. + """ + if not frozen: + return list(p.inputs) + return [n for n in p.inputs if n.name not in frozen] + + +def extract_subgraph( + gm: GraphModule, p: Partition, frozen: dict[str, torch.Tensor] | None = None +) -> GraphModule: + """Build a standalone GraphModule for one partition. + + The result takes the partition's boundary inputs as placeholders and + returns its boundary outputs, so it can be exported to ONNX on its own. + + Boundary inputs named in `frozen` become module attributes instead of + placeholders. AOTAutograd lifts every parameter and buffer to a graph + input, so without this a 2-conv block crosses the boundary with 13 tensors + instead of 1: each one is copied to the device on every call, and the ONNX + exporter sees weights as graph inputs rather than initializers, which stops + the QDQ pass from folding them to int8. Both effects are large enough to + make the offloaded graph slower than eager. + """ + frozen = frozen or {} + new_graph = Graph() + env: dict[Node, Node] = {} + consts: dict[str, torch.Tensor] = {} + + # Placeholders first, so the signature matches runtime_inputs() order. + for inp in p.inputs: + if inp.name in frozen: + continue + ph = new_graph.placeholder(inp.name) + ph.meta = dict(inp.meta) + env[inp] = ph + + for inp in p.inputs: + if inp.name not in frozen: + continue + attr = f"_frozen_{inp.name}" + consts[attr] = frozen[inp.name] + node = new_graph.get_attr(attr) + node.meta = dict(inp.meta) + env[inp] = node + + for node in p.nodes: + env[node] = new_graph.node_copy(node, lambda n: env[n]) + + new_graph.output(tuple(env[o] for o in p.outputs)) + new_graph.lint() + + # GraphModule resolves every get_attr target against the root at + # construction time, so the constants have to be on `gm` before the call, + # not registered on the result afterwards. They are removed again below to + # leave the caller's module as it was. + # `detach` runs under the backend's active FakeTensorMode, which rejects + # real tensors, so it has to be suspended (as with the ONNX export). + from torch._subclasses.fake_tensor import unset_fake_temporarily + + with unset_fake_temporarily(): + for attr, t in consts.items(): + setattr(gm, attr, t.detach()) + try: + sub = GraphModule(gm, new_graph) + finally: + for attr in consts: + if hasattr(gm, attr): + delattr(gm, attr) + + sub.graph.eliminate_dead_code() + sub.recompile() + return sub + + +def summarize(gm: GraphModule, partitions: list[Partition]) -> str: + """Human-readable partition report, for logs and debugging.""" + total = sum(1 for n in gm.graph.nodes if n.op == "call_function") + offloaded = sum(1 for p in partitions for n in p.nodes if n.op == "call_function") + lines = [ + f"{len(partitions)} BPU partition(s), " + f"{offloaded}/{total} compute nodes offloaded" + ] + for i, p in enumerate(partitions): + ops = [ + str(n.target).split(".")[-2] if "." in str(n.target) else str(n.target) + for n in p.nodes + if n.op == "call_function" + ] + lines.append( + f" [{i}] {len(p)} nodes, " + f"{len(p.inputs)} in / {len(p.outputs)} out: {' '.join(ops[:8])}" + + (" ..." if len(ops) > 8 else "") + ) + return "\n".join(lines) diff --git a/torch_fl/accelerator/bpu/qdq.py b/torch_fl/accelerator/bpu/qdq.py new file mode 100644 index 00000000..7a8e1c6e --- /dev/null +++ b/torch_fl/accelerator/bpu/qdq.py @@ -0,0 +1,142 @@ +# 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. + +"""Insert QuantizeLinear/DequantizeLinear pairs into a float ONNX graph. + +hbdk4 lowers a conv to the BPU only when its input type is si8/si16. Wrapping +each eligible op's inputs in a Q/DQ pair is what communicates that, and the +compiler's own frontend maps those to `qnt.quantize`. + +Constraints imposed by hbdk4's frontend (see onnx/opset13.py::QuantizeLinear): +scale and zero_point must be constant initializers, and the target type must be +signed. Violating any of them raises at export time, not compile time. +""" + +from __future__ import annotations + +import logging + +import numpy as np + +log = logging.getLogger("torch_fl.bpu") + +# Ops worth quantizing: the ones that actually run on the MAC array. Quantizing +# anything else only adds rescale nodes without moving work onto the BPU. +_QUANTIZABLE = {"Conv", "Gemm", "MatMul"} + + +def _sym_scale(arr: np.ndarray) -> float: + m = float(np.abs(arr).max()) if arr.size else 0.0 + return (m / 127.0) if m > 0 else 1.0 + + +def quantize_onnx( + proto, + act_scales: dict[str, float] | None = None, + default_act_scale: float = 0.05, +): + """Return a copy of `proto` with Q/DQ inserted around quantizable ops. + + Args: + proto: a float ONNX ModelProto (already shape-inferred). + act_scales: tensor-name -> scale, from calibration. Names not present + fall back to `default_act_scale`. + default_act_scale: used for activations with no calibration entry. + + Weights are quantized per-tensor from their own values, which needs no + calibration data — only activations do. + """ + import onnx + from onnx import helper, numpy_helper + + act_scales = act_scales or {} + g = proto.graph + inits = {i.name: i for i in g.initializer} + new_inits = [] + new_nodes = [] + uid = [0] + + def fresh(stem: str) -> str: + uid[0] += 1 + return f"{stem}_tbpu{uid[0]}" + + def add_const(name: str, arr: np.ndarray) -> str: + new_inits.append(numpy_helper.from_array(arr, name)) + return name + + def qdq_activation(src: str, scale: float) -> str: + """Emit Q->DQ on an activation edge, returning the new edge name.""" + s = add_const(fresh(f"{src}_s"), np.array(scale, dtype=np.float32)) + z = add_const(fresh(f"{src}_z"), np.array(0, dtype=np.int8)) + q, dq = fresh(f"{src}_q"), fresh(f"{src}_dq") + new_nodes.append(helper.make_node("QuantizeLinear", [src, s, z], [q])) + new_nodes.append(helper.make_node("DequantizeLinear", [q, s, z], [dq])) + return dq + + def dq_weight(name: str) -> str | None: + """Replace a float weight initializer with int8 + DequantizeLinear.""" + init = inits.get(name) + if init is None: + return None + w = numpy_helper.to_array(init) + if w.dtype != np.float32: + return None + sc = _sym_scale(w) + wq = np.clip(np.rint(w / sc), -127, 127).astype(np.int8) + qn = add_const(fresh(f"{name}_q"), wq) + s = add_const(fresh(f"{name}_s"), np.array(sc, dtype=np.float32)) + z = add_const(fresh(f"{name}_z"), np.array(0, dtype=np.int8)) + out = fresh(f"{name}_dq") + new_nodes.append(helper.make_node("DequantizeLinear", [qn, s, z], [out])) + return out + + # Rewrite quantizable nodes in place, preserving graph order. + quantized = 0 + for node in g.node: + if node.op_type not in _QUANTIZABLE or len(node.input) < 2: + new_nodes.append(node) + continue + + args = list(node.input) + + # Input activation: Q/DQ using the calibrated scale. + act = args[0] + scale = act_scales.get(act, default_act_scale) + args[0] = qdq_activation(act, scale) + + # Weight: fold into int8 + DQ. If it isn't a constant initializer + # (rare, e.g. a computed weight) leave it float. + wq = dq_weight(args[1]) + if wq is not None: + args[1] = wq + + del node.input[:] + node.input.extend(args) + new_nodes.append(node) + quantized += 1 + + if not quantized: + log.warning("qdq: no quantizable ops found; graph unchanged") + return proto + + del g.node[:] + g.node.extend(new_nodes) + g.initializer.extend(new_inits) + + # Q/DQ changed the type of intermediate edges, so stale value_info would + # now be wrong. Drop it and re-infer. + del g.value_info[:] + proto = onnx.shape_inference.infer_shapes(proto, data_prop=True) + log.info("qdq: quantized %d op(s)", quantized) + return proto diff --git a/torch_fl/accelerator/bpu/runtime.py b/torch_fl/accelerator/bpu/runtime.py new file mode 100644 index 00000000..4b3b5296 --- /dev/null +++ b/torch_fl/accelerator/bpu/runtime.py @@ -0,0 +1,265 @@ +# 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. + +"""Tensor-level wrapper over hbm_runtime. + +Bridges torch.Tensor and the numpy-based hbm_runtime.run() interface, applying +the quantization parameters baked into the compiled artifact. + +The hbm_runtime API surface used here (verified against runtime 3.13.6 / HBRT +4.7.5): model metadata is exposed as *properties* keyed by model name, not +methods, and QuantParams carries (axis, quant_type, scale, zero_point). + +A tensor already on the flagos device does not need a device-to-host copy to +become a numpy array: UCP memory is mapped into this process, so `data_ptr()` is +host-dereferenceable and a numpy array can wrap it in place. See `_device_view`. +""" + +from __future__ import annotations + +import ctypes +import logging + +import numpy as np +import torch + +log = logging.getLogger("torch_fl.bpu") + +# torch dtype -> (ctypes scalar, numpy dtype), for wrapping device storage in a +# numpy array without copying. Only fixed-width scalar types the BPU can carry. +_CTYPES_OF = { + torch.float32: (ctypes.c_float, np.float32), + torch.float16: (ctypes.c_uint16, np.float16), # no ctypes half; reinterpret + torch.int8: (ctypes.c_int8, np.int8), + torch.uint8: (ctypes.c_uint8, np.uint8), + torch.int16: (ctypes.c_int16, np.int16), + torch.int32: (ctypes.c_int32, np.int32), + torch.int64: (ctypes.c_int64, np.int64), + torch.bool: (ctypes.c_bool, np.bool_), +} + + +def _device_view(t: torch.Tensor) -> np.ndarray | None: + """Zero-copy numpy view over a flagos tensor's storage, or None. + + UCP device memory is host-mapped (hbUCPMallocCached hands back a writable + `virAddr`), so the storage can be read and written in place rather than + copied to the host first. Returns None whenever a view would be wrong -- + a non-flagos tensor, a non-contiguous layout, or a dtype with no fixed-width + ctypes spelling -- and the caller falls back to `.cpu()`. + + The returned array borrows `t`'s storage and must not outlive it. Every + caller here uses it within one `__call__`, where the inputs are held alive by + the argument list. + """ + if t.device.type != "flagos" or not t.is_contiguous() or t.numel() == 0: + return None + spec = _CTYPES_OF.get(t.dtype) + if spec is None: + return None + ctype, np_dtype = spec + buf = (ctype * t.numel()).from_address(t.data_ptr()) + arr = np.frombuffer(buf, dtype=np_dtype, count=t.numel()) + return arr.reshape(tuple(t.shape)) + + +def _as_numpy(t: torch.Tensor) -> np.ndarray: + """Read a tensor as numpy, without a device-to-host copy where possible.""" + t = t.detach() + view = _device_view(t) + if view is not None: + return view + return t.contiguous().cpu().numpy() + + +# hbDNNDataType enum name -> numpy dtype. +_NP_DTYPE = { + "F32": np.float32, + "F16": np.float16, + "S8": np.int8, + "U8": np.uint8, + "S16": np.int16, + "U16": np.uint16, + "S32": np.int32, + "U32": np.uint32, + "S64": np.int64, + "BOOL8": np.bool_, +} + + +def _dtype_of(enum_val) -> np.dtype: + """Map an hbDNNDataType enum member to a numpy dtype.""" + name = getattr(enum_val, "name", str(enum_val).rsplit(".", 1)[-1]) + return _NP_DTYPE.get(name, np.float32) + + +def _scale_array(quant) -> np.ndarray | None: + """Extract a quantization scale as an array, or None if unquantized.""" + if quant is None: + return None + scale = getattr(quant, "scale", None) + if scale is None: + return None + arr = np.atleast_1d(np.asarray(scale, dtype=np.float32)) + if arr.size == 0 or not np.any(arr): + return None + return arr + + +def _broadcast(scale: np.ndarray, ndim: int, axis: int) -> np.ndarray: + """Shape a per-channel scale for broadcasting against an ndim tensor.""" + if scale.size == 1: + return scale.reshape(()) + shape = [1] * ndim + shape[axis if -ndim <= axis < ndim else -1] = -1 + return scale.reshape(shape) + + +class BPURuntime: + """Executes one compiled .hbm graph, taking and returning torch tensors.""" + + def __init__( + self, + hbm_path: str, + input_names: list[str] | None = None, + output_names: list[str] | None = None, + ): + from hbm_runtime import HB_HBMRuntime + + self._rt = HB_HBMRuntime(str(hbm_path)) + self._hbm_path = str(hbm_path) + self._model = self._rt.model_names[0] + m = self._model + + # Trust the artifact's own IO names over the ONNX export names: hbdk4 + # may rename or reorder during compilation. + self._in_names: list[str] = list(self._rt.input_names[m]) + self._out_names: list[str] = list(self._rt.output_names[m]) + self._in_dtypes = self._rt.input_dtypes[m] + self._out_dtypes = self._rt.output_dtypes[m] + self._in_shapes = self._rt.input_shapes[m] + self._in_quants = self._rt.input_quants[m] + self._out_quants = self._rt.output_quants[m] + + if input_names is not None and len(input_names) != len(self._in_names): + log.warning( + "%s: artifact takes %d input(s) %s, graph supplies %d", + hbm_path, + len(self._in_names), + self._in_names, + len(input_names), + ) + + # -- conversion ------------------------------------------------------ + + def _to_device(self, name: str, t: torch.Tensor) -> np.ndarray: + """Convert a torch tensor into the artifact's expected array. + + A flagos tensor whose dtype already matches the artifact needs no work at + all: the array returned by `_as_numpy` *is* its UCP storage. Anything + that changes dtype or scale (the quantization branch below, or the final + astype) copies, which is unavoidable -- the artifact wants int8 where the + graph holds float32. + """ + arr = _as_numpy(t) + target = _dtype_of(self._in_dtypes.get(name)) + + scale = _scale_array(self._in_quants.get(name)) + if scale is not None and np.issubdtype(target, np.integer): + quant = self._in_quants[name] + axis = getattr(quant, "axis", 1) or 1 + zp = getattr(quant, "zero_point", None) + arr = arr.astype(np.float32) / _broadcast(scale, arr.ndim, axis) + if zp is not None: + zp_arr = np.atleast_1d(np.asarray(zp, dtype=np.float32)) + if np.any(zp_arr): + arr = arr + _broadcast(zp_arr, arr.ndim, axis) + info = np.iinfo(target) + arr = np.clip(np.rint(arr), info.min, info.max) + + # asarray, not astype: astype always copies, so a zero-copy device view + # whose dtype already matches would be duplicated for nothing. + return np.ascontiguousarray(np.asarray(arr, dtype=target)) + + def _from_device( + self, name: str, arr: np.ndarray, device: torch.device | None = None + ) -> torch.Tensor: + """Dequantize an output array back to a float tensor on `device`.""" + out = arr.astype(np.float32) + scale = _scale_array(self._out_quants.get(name)) + if scale is not None and np.issubdtype(arr.dtype, np.integer): + quant = self._out_quants[name] + axis = getattr(quant, "axis", 1) or 1 + zp = getattr(quant, "zero_point", None) + if zp is not None: + zp_arr = np.atleast_1d(np.asarray(zp, dtype=np.float32)) + if np.any(zp_arr): + out = out - _broadcast(zp_arr, out.ndim, axis) + out = out * _broadcast(scale, out.ndim, axis) + t = torch.from_numpy(np.ascontiguousarray(out)) + # Outputs follow the inputs' device, so a partition spliced into a graph + # of flagos tensors hands flagos tensors to its successors rather than + # forcing the rest of the graph onto the CPU. + return t if device is None or device.type == "cpu" else t.to(device) + + # -- execution ------------------------------------------------------- + + def __call__(self, *inputs: torch.Tensor) -> list[torch.Tensor]: + if len(inputs) != len(self._in_names): + raise ValueError( + f"{self._hbm_path}: expected {len(self._in_names)} input(s) " + f"{self._in_names}, got {len(inputs)}" + ) + feed = { + name: self._to_device(name, t) for name, t in zip(self._in_names, inputs) + } + result = self._rt.run(feed) + out_map = result[self._model] + device = inputs[0].device if inputs else None + return [self._from_device(n, out_map[n], device) for n in self._out_names] + + # -- introspection --------------------------------------------------- + + @property + def input_names(self) -> list[str]: + return list(self._in_names) + + @property + def output_names(self) -> list[str]: + return list(self._out_names) + + def describe(self) -> str: + """Multi-line summary of the artifact's IO signature.""" + lines = [f"{self._hbm_path} [{self._model}]"] + for n in self._in_names: + q = _scale_array(self._in_quants.get(n)) + lines.append( + f" in {n}: {list(self._in_shapes.get(n, []))} " + f"{_dtype_of(self._in_dtypes.get(n)).__name__}" + + (f" scale[{q.size}]" if q is not None else "") + ) + for n in self._out_names: + q = _scale_array(self._out_quants.get(n)) + lines.append( + f" out {n}: " + f"{_dtype_of(self._out_dtypes.get(n)).__name__}" + + (f" scale[{q.size}]" if q is not None else "") + ) + return "\n".join(lines) + + def __repr__(self) -> str: + return ( + f"BPURuntime({self._hbm_path!r}, in={self._in_names}, " + f"out={self._out_names})" + ) diff --git a/torch_fl/configs/backends_bpu.conf b/torch_fl/configs/backends_bpu.conf new file mode 100644 index 00000000..7cce882e --- /dev/null +++ b/torch_fl/configs/backends_bpu.conf @@ -0,0 +1,17 @@ +# flagos op backend config for D-Robotics RDK BPU +# Format: op_name = backend +# Override: FLAGOS_BACKEND_CONFIG=/path/to/backends.conf +# +# Intentionally empty. The table this file fills is consulted by +# GetBackendForOp() only from kernels registered on PrivateUse1, and BPU +# registers none: the BPU's unit of execution is a whole compiled graph (a .hbm +# produced by hbdk4), not an individual operator, so there is nothing to +# dispatch per call. Every aten op reaches cpu_fallback and acceleration comes +# from the torch.compile backend instead: +# +# import torch, torch_fl +# compiled = torch.compile(model, backend="bpu") +# +# The file exists so csrc/aten/common.cc's platform lookup +# (backends_bpu.conf) resolves rather than falling through to +# backends.conf, which would log a misleading CUDA-oriented table.