diff --git a/.flake8 b/.flake8 index 5ec52507..fc4155ff 100644 --- a/.flake8 +++ b/.flake8 @@ -8,9 +8,12 @@ exclude = .git, __pycache__, build, + **/build, env, + .venv, dist, deps, + .llvm-project, *docs/build, *docs/source/generated, *docs/source/stubs, diff --git a/.github/workflows/ci-cpu.yml b/.github/workflows/ci-cpu.yml new file mode 100644 index 00000000..bcf4bbfe --- /dev/null +++ b/.github/workflows/ci-cpu.yml @@ -0,0 +1,123 @@ +# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +name: CPU CI + +on: + push: + branches: [main] + pull_request: + branches: [main] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + lint: + name: Lint + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - name: Checkout repository + uses: actions/checkout@v6 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.12" + + - name: Install lint dependencies + run: python -m pip install flake8 + + - name: Run flake8 + run: flake8 + + - name: Run cpplint + run: python scripts/cpplint.py + + - name: Check license headers + run: python scripts/check_license.py + + - name: Check inline samples are up to date + run: python test/tools/inline_samples.py --check + + build: + name: Build CPU Wheel (Python 3.12) + runs-on: ubuntu-24.04 + timeout-minutes: 360 + env: + CUTILE_BUILD_VARIANT: cpu + steps: + - name: Checkout repository + uses: actions/checkout@v6 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.12" + + - name: Install native build dependencies + run: | + sudo apt-get update + sudo apt-get install --yes clang lld g++ + + - name: Set up uv + uses: astral-sh/setup-uv@v7 + with: + enable-cache: true + + - name: Install CPU build dependencies + run: uv sync --group cpu --no-install-project + + - name: Build CPU wheel + run: uv build --wheel --no-build-isolation --python .venv/bin/python + + - name: Upload CPU wheel + uses: actions/upload-artifact@v6 + with: + name: wheel-cpu-py3.12-linux-x86_64 + path: dist/*.whl + if-no-files-found: error + retention-days: 7 + + test-cpu: + name: CPU Tests (Python 3.12) + needs: build + runs-on: ubuntu-24.04 + timeout-minutes: 360 + steps: + - name: Checkout repository + uses: actions/checkout@v6 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.12" + + - name: Download CPU wheel + uses: actions/download-artifact@v7 + with: + name: wheel-cpu-py3.12-linux-x86_64 + path: dist/ + + - name: Install test dependencies + run: | + python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + python -m pip install dist/*.whl numpy pytest pytest-benchmark pytest-env + python -m pip install setuptools==84.0.0 'cmake>=3.24,<4' ninja + + - name: Run CPU backend tests + run: python -m pytest --durations=10 --junitxml=test-results.xml -m "not benchmark" backends/cpu backends/test + + - name: Upload test results + if: ${{ !cancelled() }} + uses: actions/upload-artifact@v6 + with: + name: cpu-test-results + path: test-results.xml + retention-days: 7 \ No newline at end of file diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d5a3e892..988b7b58 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -5,11 +5,11 @@ name: CI on: - push: - branches: [main] - pull_request: - branches: [main] workflow_dispatch: + # push: + # branches: [main] + # pull_request: + # branches: [main] concurrency: group: ${{ github.workflow }}-${{ github.ref }} diff --git a/.gitignore b/.gitignore index 9c7d5afe..b7e5a99a 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,7 @@ deps *.nsys-rep *.so build/ +.llvm-project/ .cache .idea .rgignore diff --git a/CMakeLists.txt b/CMakeLists.txt index 2ca5553f..1df6311e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -114,6 +114,11 @@ endif() add_subdirectory(cext) add_subdirectory(experimental/cuda-lang/) +# The non-CUDA variant (cpu) additionally builds its backend; see backends/README.md. +if (CUTILE_BUILD_VARIANT AND NOT CUTILE_BUILD_VARIANT STREQUAL "cuda") + add_subdirectory(backends) +endif() + option(DISABLE_INTERNAL "disable internal" OFF) if (NOT DISABLE_INTERNAL) diff --git a/README.md b/README.md index 8661df7e..c62c38d2 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,46 @@ cuTile Python is a programming language for NVIDIA GPUs. The official documentat on [docs.nvidia.com](https://docs.nvidia.com/cuda/cutile-python), or built from source located in the [docs](docs/) folder. +> [!NOTE] +> Non-CUDA Backends (Experimental) +> -------------------------------- +> +> Non-CUDA backend support on Linux is experimental. The CPU backend is available, +> and Intel GPU (XPU) support is planned. Backends integrate by intercepting +> TileIR, a design that could also support other platforms. +> +> To build the CPU backend from source, install the prerequisites below, then run +> these commands from the repository root: +> +> ```bash +> uv venv --python 3.12 +> source .venv/bin/activate +> +> CUTILE_BUILD_VARIANT=cpu uv sync --group cpu --no-build-isolation-package cuda-tile +> ``` +> +> - **CPU:** requires GCC with OpenMP, Clang, and LLD. See the +> [CPU guide](backends/cpu/README.md) for prerequisites, sample runs, and usage. +> +> The CPU build does not require an NVIDIA GPU or CUDA driver. Each environment supports +> one build variant. See the [backend overview](backends/README.md) for runtime +> selection and shared behavior. +> +> Autotuning (New Feature) +>------------------------ +> +> `ct.autotune` benchmarks candidate compile-time configurations +> on the first launch, then reuses the fastest configuration for inputs with the same layout. +> +> ```python +> @ct.autotune( +> configs=[ct.tune.Config({"TILE": 128}), ct.tune.Config({"TILE": 256})], +> grid=lambda meta: (ct.cdiv(meta["N"], meta["TILE"]),), +> ) +> @ct.kernel +> def my_kernel(x, N: ct.Constant[int], TILE: ct.Constant[int]): +> ... +> ``` Example ------- diff --git a/backends/CMakeLists.txt b/backends/CMakeLists.txt new file mode 100644 index 00000000..316d2b80 --- /dev/null +++ b/backends/CMakeLists.txt @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# LLVM builds the compiler as an external project, separately from the launchers. +if(TARGET TileIRToMLIR) + include("${CMAKE_CURRENT_SOURCE_DIR}/compiler.cmake") + return() +endif() + +# Builds the native parts of the non-CUDA backend selected by CUTILE_BUILD_VARIANT: +# the native compiler and the backend's launcher extension. Installing +# the build (`cmake --install --prefix --component cuTileBackend`) +# stages them and their runtime files into /cuda/tile/_backends, next to the +# Python backend in src/cuda/tile/_backends. + +if(NOT CMAKE_SYSTEM_NAME STREQUAL "Linux") + message(FATAL_ERROR "The native CPU compiler currently supports Linux only") +endif() + +if (NOT CUTILE_BUILD_VARIANT STREQUAL "cpu") + message(FATAL_ERROR "CUTILE_BUILD_VARIANT must be one of: cuda, cpu") +endif() + +# Installing only this component skips the install rules of cuTile's other +# dependencies, such as dlpack. +set(CMAKE_INSTALL_DEFAULT_COMPONENT_NAME cuTileBackend) + +find_package(Python REQUIRED COMPONENTS Interpreter Development.Module) +execute_process( + COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir + OUTPUT_VARIABLE nanobind_ROOT + OUTPUT_STRIP_TRAILING_WHITESPACE + COMMAND_ERROR_IS_FATAL ANY) +find_package(nanobind CONFIG REQUIRED) + +# Package directory of the backends, relative to the install prefix. +set(CUTILE_BACKEND_INSTALL_DIR cuda/tile/_backends) + +include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/FetchTileIRToMLIR.cmake") +fetch_tileir_to_mlir() +include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/TileIRToMLIR.cmake") + +install(FILES "${CUTILE_COMPILER_PATH}" DESTINATION "${CUTILE_BACKEND_INSTALL_DIR}") +install(FILES "${CUTILE_LLVM_LICENSE_PATH}" + DESTINATION "${CUTILE_BACKEND_INSTALL_DIR}/_licenses/llvm") +install(FILES "${CUTILE_TILEIR_TO_MLIR_DIR}/LICENSE" + DESTINATION "${CUTILE_BACKEND_INSTALL_DIR}/_licenses/tileir-to-mlir") +install(FILES "${CUTILE_CUDA_TILE_SOURCE_DIR}/LICENSE.txt" + DESTINATION "${CUTILE_BACKEND_INSTALL_DIR}/_licenses/cuda-tile") + +add_subdirectory("${CUTILE_BUILD_VARIANT}") diff --git a/backends/README.md b/backends/README.md new file mode 100644 index 00000000..2c72d5fb --- /dev/null +++ b/backends/README.md @@ -0,0 +1,43 @@ + + + +# Non-CUDA backends + +The built-in CPU backend runs cuTile kernels on Linux host arrays and CPU +PyTorch tensors. It is experimental; see the [CPU guide](cpu/README.md) for +requirements, installation, and a sample. + +## Select a backend + +```python +import cuda.tile as ct + +ct.set_backend("cpu") +ct.launch(None, grid, kernel, args) +ct.set_backend(None) # Restore the CUDA launch path. +``` + +Backend selection applies to the whole process and is not inferred from the +arguments. Coordinate changes between threads. CPU launches are synchronous; +pass CPU arrays and `None` for the stream. CUDA-only launch options are not +available on the CPU path. + +The CPU backend must be installed in the environment before selecting it. +Runtime selection does not change the installed compiler or build variant. + +## Options and limits + +Use `ct.backend_options({...})` to scope backend options to a block of code. +Options apply to the current thread, and unknown options are rejected. The CPU +backend supports `num_cpu_threads` and `assume_in_bounds`; see the +[CPU guide](cpu/README.md#use-in-a-program). + +Compiled kernels are cached. Changing a code-generation option selects a +different compiled kernel; changing the CPU thread count does not recompile. +CUDA-specific hints such as `num_ctas` and `occupancy` do not apply. + +## Tests + +Run the backend tests from the repository root with `python -m pytest backends`. +The CPU-specific guide has instructions for installing the test +dependencies and running only CPU tests. diff --git a/backends/cmake/BuildLLVM.cmake.in b/backends/cmake/BuildLLVM.cmake.in new file mode 100644 index 00000000..2ee18675 --- /dev/null +++ b/backends/cmake/BuildLLVM.cmake.in @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +cmake_minimum_required(VERSION 3.20) +set(_source [==[@CUTILE_LLVM_SOURCE_DIR@/llvm]==]) +set(_binary [==[@CUTILE_LLVM_BINARY_DIR@]==]) +set(_args [==[@_cutile_llvm_cmake_args@]==]) +set(_targets [==[@_cutile_llvm_build_targets@]==]) +string(REPLACE "|" "\\;" _args "${_args}") +string(SHA256 _request "${CMAKE_COMMAND};${_source};${_args}") + +set(_cache "${_binary}/CMakeCache.txt") +set(_stamp "${_binary}/cutile-configure.stamp") +set(_state "") +if(EXISTS "${_cache}" AND EXISTS "${_binary}/build.ninja") + file(SHA256 "${_cache}" _cache_hash) + set(_state "${_request}\n${_cache_hash}\n") +endif() +set(_saved "") +if(EXISTS "${_stamp}") + file(READ "${_stamp}" _saved) +endif() +# The shared LLVM tree may have been reconfigured by another outer build or +# manually. Ninja handles source/CMake-file changes on otherwise matching builds. +if(NOT _state OR NOT _state STREQUAL _saved) + file(REMOVE "${_stamp}") + execute_process( + COMMAND "${CMAKE_COMMAND}" -S "${_source}" -B "${_binary}" -G Ninja ${_args} + COMMAND_ERROR_IS_FATAL ANY) +endif() +execute_process( + COMMAND "${CMAKE_COMMAND}" --build "${_binary}" --target ${_targets} + COMMAND_ERROR_IS_FATAL ANY) +file(SHA256 "${_cache}" _cache_hash) +file(WRITE "${_stamp}" "${_request}\n${_cache_hash}\n") diff --git a/backends/cmake/CheckoutLLVM.cmake b/backends/cmake/CheckoutLLVM.cmake new file mode 100644 index 00000000..d9dd87d5 --- /dev/null +++ b/backends/cmake/CheckoutLLVM.cmake @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +cmake_minimum_required(VERSION 3.19) + +foreach(_required GIT_EXECUTABLE SOURCE_DIR TILEIR_REQUIRED_LLVM_REVISION) + if(NOT DEFINED ${_required} OR "${${_required}}" STREQUAL "") + message(FATAL_ERROR "${_required} is required") + endif() +endforeach() + +if(NOT EXISTS "${SOURCE_DIR}/.git") + if(EXISTS "${SOURCE_DIR}") + file(GLOB _source_contents "${SOURCE_DIR}/*" "${SOURCE_DIR}/.[!.]*") + if(_source_contents) + message(FATAL_ERROR + "LLVM source directory exists but is not a Git checkout: ${SOURCE_DIR}") + endif() + endif() + execute_process( + COMMAND "${GIT_EXECUTABLE}" clone --filter=blob:none --no-checkout + https://github.com/llvm/llvm-project.git "${SOURCE_DIR}" + COMMAND_ERROR_IS_FATAL ANY) +endif() + +set(_required_revisions "${TILEIR_REQUIRED_LLVM_REVISION}") +if(DEFINED REVISION AND NOT REVISION STREQUAL "") + list(APPEND _required_revisions "${REVISION}") +endif() + +foreach(_revision IN LISTS _required_revisions) + execute_process( + COMMAND "${GIT_EXECUTABLE}" -C "${SOURCE_DIR}" cat-file -e "${_revision}^{commit}" + RESULT_VARIABLE _revision_exists + ERROR_QUIET) + if(NOT _revision_exists EQUAL 0) + execute_process( + COMMAND "${GIT_EXECUTABLE}" -C "${SOURCE_DIR}" fetch --filter=blob:none + origin "${_revision}" + COMMAND_ERROR_IS_FATAL ANY) + endif() +endforeach() + +set(_selected_revision "${TILEIR_REQUIRED_LLVM_REVISION}") +if(DEFINED REVISION AND NOT REVISION STREQUAL "") + execute_process( + COMMAND "${GIT_EXECUTABLE}" -C "${SOURCE_DIR}" merge-base --is-ancestor + "${_selected_revision}" "${REVISION}" + RESULT_VARIABLE _required_is_ancestor) + if(NOT _required_is_ancestor EQUAL 0) + message(FATAL_ERROR + "LLVM override ${REVISION} does not descend from required revision " + "${_selected_revision}") + endif() + set(_selected_revision "${REVISION}") +endif() + +execute_process( + COMMAND "${GIT_EXECUTABLE}" -C "${SOURCE_DIR}" rev-parse HEAD + RESULT_VARIABLE _head_result + OUTPUT_VARIABLE _head_revision + OUTPUT_STRIP_TRAILING_WHITESPACE + ERROR_QUIET) +if(_head_result EQUAL 0 AND _head_revision STREQUAL _selected_revision) + return() +endif() + +execute_process( + COMMAND "${GIT_EXECUTABLE}" -C "${SOURCE_DIR}" checkout --detach "${_selected_revision}" + COMMAND_ERROR_IS_FATAL ANY) diff --git a/backends/cmake/FetchTileIRToMLIR.cmake b/backends/cmake/FetchTileIRToMLIR.cmake new file mode 100644 index 00000000..9ed4b258 --- /dev/null +++ b/backends/cmake/FetchTileIRToMLIR.cmake @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +set(TILEIR_TO_MLIR_REPOSITORY "https://github.com/libxsmm/tileir-to-mlir" + CACHE STRING "TileIRToMLIR repo URL") +set(TILEIR_TO_MLIR_REVISION "ba6c9e50274522089b7ad99b3e8a571480481330" + CACHE STRING "TileIRToMLIR revision") +set(TILEIR_TO_MLIR_REVISION_SHA256 + "f39644ec0eb4d98db8de6b8bcf51f82d67ea80081080e6c397695a20369eb0d4" + CACHE STRING "Expected SHA256 of the pinned TileIRToMLIR archive") + +set(CUDA_TILE_REPOSITORY "https://github.com/NVIDIA/cuda-tile" + CACHE STRING "CUDA Tile repo URL") +set(CUDA_TILE_REVISION "e11729ed4e228af005dc8cc113d7b84157e9d654" + CACHE STRING "CUDA Tile revision required by TileIRToMLIR") +set(CUDA_TILE_REVISION_SHA256 + "56c03d732ac5f23e0ae9c0c184ff65aa282c24bd30a1bde723457f8e986e855e" + CACHE STRING "Expected SHA256 of the pinned CUDA Tile archive") + +macro(fetch_tileir_to_mlir) + include(FetchContent) + message(STATUS "Fetching TileIRToMLIR: ${TILEIR_TO_MLIR_REVISION}") + FetchContent_Declare( + tileir_to_mlir + URL "${TILEIR_TO_MLIR_REPOSITORY}/archive/${TILEIR_TO_MLIR_REVISION}.tar.gz" + URL_HASH "SHA256=${TILEIR_TO_MLIR_REVISION_SHA256}" + SOURCE_SUBDIR does-not-exist + DOWNLOAD_EXTRACT_TIMESTAMP TRUE) + message(STATUS "Fetching CUDA Tile: ${CUDA_TILE_REVISION}") + FetchContent_Declare( + cuda_tile + URL "${CUDA_TILE_REPOSITORY}/archive/${CUDA_TILE_REVISION}.tar.gz" + URL_HASH "SHA256=${CUDA_TILE_REVISION_SHA256}" + SOURCE_SUBDIR does-not-exist + DOWNLOAD_EXTRACT_TIMESTAMP TRUE) + FetchContent_MakeAvailable(tileir_to_mlir cuda_tile) + set(CUTILE_TILEIR_TO_MLIR_DIR "${tileir_to_mlir_SOURCE_DIR}") + set(CUTILE_CUDA_TILE_SOURCE_DIR "${cuda_tile_SOURCE_DIR}") +endmacro() \ No newline at end of file diff --git a/backends/cmake/TileIRToMLIR.cmake b/backends/cmake/TileIRToMLIR.cmake new file mode 100644 index 00000000..a40e62f7 --- /dev/null +++ b/backends/cmake/TileIRToMLIR.cmake @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +set(CUTILE_TOOLCHAIN_DIR "" CACHE PATH + "Existing native compiler installation for CUTILE_BUILD_VARIANT; skips building LLVM") + +execute_process( + COMMAND "${Python_EXECUTABLE}" -c + "import sysconfig; print(sysconfig.get_config_var('EXT_SUFFIX'))" + OUTPUT_VARIABLE _compiler_suffix OUTPUT_STRIP_TRAILING_WHITESPACE + COMMAND_ERROR_IS_FATAL ANY) +set(_compiler_name "_compiler${_compiler_suffix}") + +set(_tileir_to_mlir_dir "${CUTILE_TILEIR_TO_MLIR_DIR}") +if(NOT EXISTS "${_tileir_to_mlir_dir}/CMakeLists.txt") + message(FATAL_ERROR + "TileIRToMLIR sources are missing: ${_tileir_to_mlir_dir}") +endif() + +if(CUTILE_TOOLCHAIN_DIR) + set(CUTILE_COMPILER_PATH "${CUTILE_TOOLCHAIN_DIR}/lib/${_compiler_name}") + set(CUTILE_LLVM_LICENSE_PATH "${CUTILE_TOOLCHAIN_DIR}/share/cutile/llvm/LICENSE.TXT") + foreach(_path IN ITEMS CUTILE_COMPILER_PATH CUTILE_LLVM_LICENSE_PATH) + if(NOT EXISTS "${${_path}}") + message(FATAL_ERROR "CUTILE_TOOLCHAIN_DIR must contain ${${_path}}") + endif() + endforeach() + return() +endif() + +set(_tileir_required_llvm_revision_file + "${_tileir_to_mlir_dir}/llvm-revision.txt") +set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS + "${_tileir_required_llvm_revision_file}") +file(STRINGS "${_tileir_required_llvm_revision_file}" + _tileir_required_llvm_revision LIMIT_COUNT 1) +set(_cutile_llvm_override "") +if(DEFINED ENV{USE_LLVM_REVISION} AND NOT "$ENV{USE_LLVM_REVISION}" STREQUAL "") + set(_cutile_llvm_override "$ENV{USE_LLVM_REVISION}") +elseif(DEFINED USE_LLVM_REVISION AND NOT USE_LLVM_REVISION STREQUAL "") + set(_cutile_llvm_override "${USE_LLVM_REVISION}") +endif() + +set(CUTILE_LLVM_SOURCE_DIR "${CMAKE_SOURCE_DIR}/.llvm-project" CACHE PATH + "Persistent LLVM source checkout") +set(CUTILE_LLVM_LICENSE_PATH "${CUTILE_LLVM_SOURCE_DIR}/llvm/LICENSE.TXT") +set(CUTILE_LLVM_BINARY_DIR "${CMAKE_SOURCE_DIR}/build/llvm" CACHE PATH + "LLVM/MLIR build directory") +set(CUTILE_LLVM_BUILD_TYPE "Release" CACHE STRING + "Build type for LLVM, MLIR, and TileIRToMLIR") +# Assertions abort the process that compiles in-process; enable them only to debug. +option(CUTILE_LLVM_ENABLE_ASSERTIONS "Enable LLVM/MLIR assertions" OFF) + +find_package(Git REQUIRED) +find_program(CUTILE_CLANG_EXECUTABLE NAMES clang REQUIRED) +find_program(CUTILE_CLANGXX_EXECUTABLE NAMES clang++ REQUIRED) +find_program(CUTILE_LLD_EXECUTABLE NAMES ld.lld REQUIRED) + +include(ExternalProject) +set(_required_llvm_revisions_file + "${CMAKE_BINARY_DIR}/tileir-toolchain/required-llvm-revisions.txt") +file(CONFIGURE OUTPUT "${_required_llvm_revisions_file}" + CONTENT "tileir=${_tileir_required_llvm_revision}\noverride=${_cutile_llvm_override}\n" @ONLY) + +set(CUTILE_COMPILER_PATH "${CUTILE_LLVM_BINARY_DIR}/lib/${_compiler_name}") + +set(_cutile_checkout_script "${CMAKE_CURRENT_LIST_DIR}/CheckoutLLVM.cmake") +set(_cutile_checkout_command + "${CMAKE_COMMAND}" + "-DGIT_EXECUTABLE=${GIT_EXECUTABLE}" + "-DSOURCE_DIR=${CUTILE_LLVM_SOURCE_DIR}" + "-DTILEIR_REQUIRED_LLVM_REVISION=${_tileir_required_llvm_revision}" + "-DREVISION=${_cutile_llvm_override}" + -P "${_cutile_checkout_script}") + +include("${CMAKE_CURRENT_SOURCE_DIR}/${CUTILE_BUILD_VARIANT}/toolchain.cmake") +set(_cutile_llvm_build_targets _compiler) +set(_cutile_llvm_byproducts "${CUTILE_COMPILER_PATH}") + +set(_cutile_llvm_cmake_args + "-DCMAKE_BUILD_TYPE:STRING=${CUTILE_LLVM_BUILD_TYPE}" + "-DCMAKE_C_COMPILER:FILEPATH=${CUTILE_CLANG_EXECUTABLE}" + "-DCMAKE_CXX_COMPILER:FILEPATH=${CUTILE_CLANGXX_EXECUTABLE}" + "-DCMAKE_LINKER:FILEPATH=${CUTILE_LLD_EXECUTABLE}" + "-DLLVM_ENABLE_PROJECTS:STRING=${_cutile_llvm_projects}" + "-DLLVM_EXTERNAL_PROJECTS:STRING=TileIRToMLIR|cuTileCompiler" + "-DLLVM_EXTERNAL_TILEIRTOMLIR_SOURCE_DIR:PATH=${_tileir_to_mlir_dir}" + # Explicit, so that a shared LLVM tree never keeps another checkout's path. + "-DTILEIR_TO_MLIR_CUDA_TILE_SOURCE_DIR:PATH=${CUTILE_CUDA_TILE_SOURCE_DIR}" + "-DLLVM_EXTERNAL_CUTILECOMPILER_SOURCE_DIR:PATH=${CMAKE_CURRENT_SOURCE_DIR}" + "-DCUTILE_BUILD_VARIANT:STRING=${CUTILE_BUILD_VARIANT}" + "-DUSE_LLVM_REVISION:STRING=${_cutile_llvm_override}" + -DLLVM_USE_LINKER:STRING=lld + "-DLLVM_TARGETS_TO_BUILD:STRING=${_cutile_llvm_targets}" + "-DLLVM_EXPERIMENTAL_TARGETS_TO_BUILD:STRING=${_cutile_llvm_experimental_targets}" + "-DLLVM_ENABLE_ASSERTIONS:BOOL=${CUTILE_LLVM_ENABLE_ASSERTIONS}" + -DLLVM_INSTALL_UTILS:BOOL=ON + -DLLVM_INCLUDE_TESTS:BOOL=OFF + -DLLVM_INCLUDE_BENCHMARKS:BOOL=OFF + -DLLVM_INCLUDE_EXAMPLES:BOOL=OFF + -DLLVM_BUILD_LLVM_DYLIB:BOOL=OFF + -DLLVM_LINK_LLVM_DYLIB:BOOL=OFF + # Optional host libraries add runtime dependencies the compiler never uses. + -DLLVM_ENABLE_ZLIB:STRING=OFF + -DLLVM_ENABLE_ZSTD:STRING=OFF + -DLLVM_ENABLE_LZMA:STRING=OFF + -DLLVM_ENABLE_LIBXML2:STRING=OFF + "-DCMAKE_DISABLE_FIND_PACKAGE_ocloc:BOOL=${_cutile_disable_ocloc}" + # Clear results that the shared LLVM cache would otherwise retain from + # another variant. + -UOCLOC_FOUND + -UCMAKE_PROJECT_INCLUDE + -DMLIR_INCLUDE_TESTS:BOOL=OFF + -DMLIR_ENABLE_EXECUTION_ENGINE:BOOL=ON + -DMLIR_ENABLE_LEVELZERO_RUNNER:BOOL=OFF + -DMLIR_ENABLE_BINDINGS_PYTHON:BOOL=OFF + "-DPython3_EXECUTABLE:FILEPATH=${Python_EXECUTABLE}" + "-DPython_EXECUTABLE:FILEPATH=${Python_EXECUTABLE}" + "-Dnanobind_DIR:PATH=${nanobind_DIR}") +list(APPEND _cutile_llvm_cmake_args ${_cutile_backend_cmake_args}) + +set(_cutile_llvm_build_script "${CMAKE_BINARY_DIR}/tileir-toolchain/BuildLLVM.cmake") +configure_file("${CMAKE_CURRENT_LIST_DIR}/BuildLLVM.cmake.in" + "${_cutile_llvm_build_script}" @ONLY) + +ExternalProject_Add(cutile_llvm + PREFIX "${CMAKE_BINARY_DIR}/tileir-toolchain/llvm" + SOURCE_DIR "${CUTILE_LLVM_SOURCE_DIR}" + BINARY_DIR "${CUTILE_LLVM_BINARY_DIR}" + DOWNLOAD_COMMAND ${_cutile_checkout_command} + UPDATE_COMMAND "" + PATCH_COMMAND "" + CONFIGURE_COMMAND "" + BUILD_COMMAND "${CMAKE_COMMAND}" -P "${_cutile_llvm_build_script}" + INSTALL_COMMAND "" + BUILD_BYPRODUCTS ${_cutile_llvm_byproducts} + BUILD_ALWAYS ON + USES_TERMINAL_DOWNLOAD TRUE + USES_TERMINAL_CONFIGURE TRUE + USES_TERMINAL_BUILD TRUE) + +ExternalProject_Add_StepDependencies(cutile_llvm download + "${_required_llvm_revisions_file}") + +if(_cutile_llvm_override) + message(STATUS "LLVM revision override: ${_cutile_llvm_override}") +else() + message(STATUS "LLVM revision: ${_tileir_required_llvm_revision}") +endif() +message(STATUS "LLVM source: ${CUTILE_LLVM_SOURCE_DIR}") +message(STATUS "Native cuTile compiler: ${CUTILE_COMPILER_PATH}") + +unset(_cutile_checkout_command) +unset(_cutile_checkout_script) +unset(_required_llvm_revisions_file) +unset(_cutile_llvm_override) +unset(_tileir_required_llvm_revision) +unset(_tileir_required_llvm_revision_file) +unset(_tileir_to_mlir_dir) diff --git a/backends/compiler.cmake b/backends/compiler.cmake new file mode 100644 index 00000000..53c3be1d --- /dev/null +++ b/backends/compiler.cmake @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +if(NOT TARGET TileIRToMLIR) + message(FATAL_ERROR "Build cuTileCompiler alongside the TileIRToMLIR LLVM external project") +endif() + +find_package(Python REQUIRED COMPONENTS Interpreter Development.Module) +find_package(nanobind CONFIG REQUIRED) + +nanobind_add_module(_compiler FREE_THREADED mlir/compiler.cpp) +target_include_directories(_compiler PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}") +target_include_directories(_compiler SYSTEM PRIVATE + "${LLVM_MAIN_SRC_DIR}/../mlir/include" + "${LLVM_BINARY_DIR}/tools/mlir/include") +target_link_libraries(_compiler PRIVATE + TileIRToMLIR CudaTileBytecodeReader + MLIRParser MLIRTransforms + MLIRAffineDialect MLIRArithDialect MLIRControlFlowDialect MLIRFuncDialect + MLIRGPUDialect MLIRIndexDialect MLIRLLVMDialect MLIRMathDialect + MLIRMemRefDialect MLIRSCFDialect MLIRUBDialect + MLIRVectorDialect MLIRX86Dialect MLIRFuncInlinerExtension + MLIRArithValueBoundsOpInterfaceImpl + MLIRArithToLLVM MLIRControlFlowToLLVM MLIRFuncToLLVM MLIRIndexToLLVM + MLIRMathToLLVM MLIRMemRefToLLVM MLIRUBToLLVM MLIRVectorToLLVM + MLIRX86Transforms + MLIRBuiltinToLLVMIRTranslation MLIRLLVMToLLVMIRTranslation + MLIRVectorTransforms) +target_sources(_compiler PRIVATE mlir/LowerVectorScan.cpp) +target_link_options(_compiler PRIVATE "LINKER:--exclude-libs,ALL") +set_target_properties(_compiler PROPERTIES + LIBRARY_OUTPUT_DIRECTORY "${LLVM_BINARY_DIR}/lib" + BUILD_WITH_INSTALL_RPATH TRUE) + +include("${CMAKE_CURRENT_SOURCE_DIR}/${CUTILE_BUILD_VARIANT}/compiler.cmake") + +if(NOT LLVM_ENABLE_RTTI) + set_source_files_properties( + mlir/LowerVectorScan.cpp + PROPERTIES COMPILE_OPTIONS "-fno-rtti") +endif() + +install(TARGETS _compiler LIBRARY DESTINATION lib COMPONENT cuTileCompiler) +install(FILES "${LLVM_MAIN_SRC_DIR}/LICENSE.TXT" + DESTINATION share/cutile/llvm COMPONENT cuTileCompiler) diff --git a/backends/conftest.py b/backends/conftest.py new file mode 100644 index 00000000..b0d1f5e9 --- /dev/null +++ b/backends/conftest.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from cuda.tile._backends import _custom +from cuda.tile._backends._signature import build_signature +from cuda.tile._cext import default_tile_context + + +@pytest.fixture(autouse=True) +def _compile_in_process(monkeypatch): + """Exercise the default in-process compilation; tests opt into a timeout worker.""" + monkeypatch.setattr(default_tile_context.config, "compiler_timeout_sec", None) + + +@pytest.fixture +def native(): + return pytest.importorskip("cuda.tile._backends._compiler") + + +@pytest.fixture +def bytecode(): + def compile(kernel, arguments, backend): + return _custom.compile_bytecode(kernel, backend, build_signature(kernel, arguments)) + + return compile diff --git a/backends/cpu/ApproximateMath.cpp b/backends/cpu/ApproximateMath.cpp new file mode 100644 index 00000000..988a172b --- /dev/null +++ b/backends/cpu/ApproximateMath.cpp @@ -0,0 +1,138 @@ +// SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights +// reserved. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "ApproximateMath.h" + +#include + +#include "mlir/Dialect/Affine/IR/AffineDialect.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Math/IR/Math.h" +#include "mlir/Dialect/Math/Transforms/Passes.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Vector/IR/VectorOps.h" +#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h" +#include "mlir/Dialect/X86/Transforms.h" +#include "mlir/Dialect/X86/X86Dialect.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Dominance.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/CSE.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Transforms/RegionUtils.h" + +namespace cutile { +namespace { + +class PrepareCPUVectorsPass + : public mlir::PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PrepareCPUVectorsPass) + + explicit PrepareCPUVectorsPass(llvm::StringRef entryPoint) + : prepareX86(entryPoint != "__transform_nano"), + prepareAMX(entryPoint == "__transform_nano_amx") {} + + void getDependentDialects(mlir::DialectRegistry ®istry) const override { + registry.insert(); + } + + void runOnOperation() override { + llvm::SmallVector functions; + getOperation().walk( + [&](mlir::func::FuncOp function) { functions.push_back(function); }); + auto apply = [&](auto populatePatterns) { + for (mlir::func::FuncOp function : functions) { + mlir::RewritePatternSet patterns(&getContext()); + populatePatterns(patterns); + mlir::GreedyRewriteConfig config; + config.setMaxIterations(mlir::GreedyRewriteConfig::kNoLimit); + config.setMaxNumRewrites(mlir::GreedyRewriteConfig::kNoLimit); + if (mlir::failed(mlir::applyPatternsGreedily( + function, std::move(patterns), config))) + return mlir::failure(); + } + return mlir::success(); + }; + if (mlir::failed(apply([](mlir::RewritePatternSet &patterns) { + mlir::vector::populateCastAwayVectorLeadingOneDimPatterns(patterns); + mlir::vector::populateDropUnitDimWithShapeCastPatterns(patterns); + }))) { + signalPassFailure(); + return; + } + if (!prepareX86) + return; + if (mlir::failed( + apply(mlir::x86::populatePrepareVectorContractLoopsPatterns))) { + signalPassFailure(); + return; + } + mlir::IRRewriter rewriter(&getContext()); + for (mlir::func::FuncOp function : functions) + for (mlir::Region ®ion : function->getRegions()) + mlir::eliminateTriviallyDeadOps(rewriter, region); + for (mlir::func::FuncOp function : functions) { + mlir::DominanceInfo dominance; + mlir::eliminateCommonSubExpressions(rewriter, dominance, function); + } + if (prepareAMX && mlir::failed(apply([](mlir::RewritePatternSet &patterns) { + mlir::x86::populateVectorContractMultiLevelUnrollForAMXPatterns( + patterns); + mlir::x86::populateVectorContractToAMXDotProductPatterns(patterns); + }))) + signalPassFailure(); + } + +private: + bool prepareX86; + bool prepareAMX; +}; + +class ApproximateMathPass + : public mlir::PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ApproximateMathPass) + + void getDependentDialects(mlir::DialectRegistry ®istry) const override { + registry.insert(); + } + + void runOnOperation() override { + auto approximate = [](llvm::StringRef name) { + return name != mlir::math::SinOp::getOperationName() && + name != mlir::math::CosOp::getOperationName(); + }; + mlir::RewritePatternSet patterns(&getContext()); + // exp2 has no approximation of its own; express it through exp. + mlir::math::populateExpansionPatterns(patterns, {"exp2"}); + mlir::populateMathF32ExpansionPatterns(patterns, approximate); + mlir::populateMathPolynomialApproximationPatterns(patterns, approximate); + if (mlir::failed( + mlir::applyPatternsGreedily(getOperation(), std::move(patterns)))) + signalPassFailure(); + } +}; + +} // namespace + +std::unique_ptr createApproximateMathPass() { + return std::make_unique(); +} + +std::unique_ptr +createPrepareCPUVectorsPass(llvm::StringRef entryPoint) { + return std::make_unique(entryPoint); +} + +} // namespace cutile diff --git a/backends/cpu/ApproximateMath.h b/backends/cpu/ApproximateMath.h new file mode 100644 index 00000000..afe3087a --- /dev/null +++ b/backends/cpu/ApproximateMath.h @@ -0,0 +1,24 @@ +// SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights +// reserved. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include + +namespace llvm { +class StringRef; +} + +namespace mlir { +class Pass; +} + +namespace cutile { +// Replaces math functions with vectorizable polynomial approximations, except +// sin and cos, whose approximation is inaccurate for large arguments. +std::unique_ptr createApproximateMathPass(); +std::unique_ptr +createPrepareCPUVectorsPass(llvm::StringRef entryPoint); +} // namespace cutile diff --git a/backends/cpu/CMakeLists.txt b/backends/cpu/CMakeLists.txt new file mode 100644 index 00000000..5a20b144 --- /dev/null +++ b/backends/cpu/CMakeLists.txt @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +find_package(OpenMP REQUIRED COMPONENTS CXX) + +nanobind_add_module(_cpu FREE_THREADED launcher_cpu.cpp) +target_link_libraries(_cpu PRIVATE OpenMP::OpenMP_CXX ${CMAKE_DL_LIBS}) + +install(TARGETS _cpu DESTINATION "${CUTILE_BACKEND_INSTALL_DIR}/cpu") diff --git a/backends/cpu/README.md b/backends/cpu/README.md new file mode 100644 index 00000000..9e70e051 --- /dev/null +++ b/backends/cpu/README.md @@ -0,0 +1,99 @@ + + + +# CPU backend + +Run cuTile kernels on Linux CPUs with NumPy-compatible arrays or CPU PyTorch +tensors. Grids can have one, two, or three dimensions. The backend is +experimental, so verify results for your workload; compile times can be long +and performance depends on the kernel and host CPU. + +## Prerequisites + +- Linux, Git, and [uv](https://docs.astral.sh/uv/getting-started/installation/). +- GCC with OpenMP development headers and runtime, plus Clang and LLD to build + LLVM/MLIR. GCC builds cuTile's extensions. + +An NVIDIA GPU, CUDA driver, or system CUDA Toolkit is not required. The build +dependencies provide the CUDA headers needed by the extension. + +## Build and install + +From the repository root: + +```bash +uv venv --python 3.12 +. .venv/bin/activate +CUTILE_BUILD_VARIANT=cpu uv sync --group cpu --no-build-isolation-package cuda-tile +``` + +The initial build compiles LLVM/MLIR and can take a while. Later builds reuse +the build tree. Keep `--no-build-isolation-package cuda-tile` when syncing this +environment. + +Running kernels does not require a compiler toolchain, LLVM executables, +development files, or a source checkout. The target system must provide the +C/C++ runtimes and GCC's OpenMP runtime (`libgomp.so.1`), which are not bundled. +On Debian or Ubuntu, install the runtime with: + +```bash +sudo apt-get install libgomp1 +``` + +## Run a sample + +With the environment active, from the repository root: + +```bash +python samples/cpu/VectorAddition-cpu.py --correctness-check +``` + +More examples and benchmarks are in [samples/cpu](../../samples/cpu). Check each +script's `--help` for supported sizes, correctness checks, and thread settings. + +## Use in a program + +Select the backend, pass CPU arrays, and use `None` as the stream: + +```python +import cuda.tile as ct + +ct.set_backend("cpu") +with ct.backend_options({"num_cpu_threads": 4}): + ct.launch(None, grid, kernel, cpu_args) +``` + +`kernel` is a decorated cuTile kernel, `grid` is its block grid, and `cpu_args` +contains its arguments. Launches are synchronous. When porting a CUDA program, +keep the kernel where supported, move its arrays to the CPU, and remove CUDA +stream and device assumptions from the host code. `ct.set_backend(None)` +restores the CUDA launch path. + +### Options + +`num_cpu_threads` sets the number of OpenMP workers. The default, `0`, uses +OpenMP's maximum and honors `OMP_NUM_THREADS`. Changing this option does not +recompile kernels. + +Bounds checks are enabled by default. Set `assume_in_bounds=True` only when +every tile load and store is in bounds. + +## Troubleshooting + +Set `MLIR_ENABLE_DUMP=1` to print generated MLIR during compilation, or set +`CUTILE_CPU_DUMP_MLIR` to a file path to write the lowered MLIR. + +If the CPU launcher cannot find `libgomp.so.1`, install GCC's OpenMP runtime. +If loading a compiled kernel fails, check that `TMPDIR` permits executable +files; kernels are loaded from temporary shared libraries. + +## Run tests + +From the repository root, install the test tools and run CPU backend tests: + +```bash +uv pip install pytest pytest-benchmark pytest-env +python -m pytest backends/cpu +``` + +Tests requiring another native build variant are skipped. diff --git a/backends/cpu/amx.h b/backends/cpu/amx.h new file mode 100644 index 00000000..5c26863a --- /dev/null +++ b/backends/cpu/amx.h @@ -0,0 +1,48 @@ +// SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights +// reserved. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#if defined(__linux__) && defined(__x86_64__) +#include +#include +#include + +#include + +#ifndef ARCH_REQ_XCOMP_PERM +#define ARCH_REQ_XCOMP_PERM 0x1023 +#endif +#ifndef ARCH_XCOMP_TILEDATA +#define ARCH_XCOMP_TILEDATA 18 +#endif +#endif + +namespace cutile::cpu { + +// Checks the required AMX features and requests Linux tile-data permission. +inline bool enableAMX() { +#if defined(__linux__) && defined(__x86_64__) + // CPUID leaf 7, subleaf 0: AMX-BF16, AMX-TILE, and AMX-INT8 are EDX bits 22, + // 24, and 25. + constexpr unsigned int amxBf16 = 1u << 22; + constexpr unsigned int amxTile = 1u << 24; + constexpr unsigned int amxInt8 = 1u << 25; + constexpr unsigned int requiredFeatures = amxBf16 | amxTile | amxInt8; + // CPUID returns four register values; AMX feature bits are in EDX. + unsigned int cpuidEax, cpuidEbx, cpuidEcx, cpuidEdx; + if (!__get_cpuid_count(7, 0, &cpuidEax, &cpuidEbx, &cpuidEcx, &cpuidEdx)) + return false; + if ((cpuidEdx & requiredFeatures) != requiredFeatures) + return false; + + // Linux requires per-thread permission for the XTILEDATA xstate component. + return syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM, ARCH_XCOMP_TILEDATA) == 0; +#else + return false; +#endif +} + +} // namespace cutile::cpu diff --git a/backends/cpu/compiler.cmake b/backends/cpu/compiler.cmake new file mode 100644 index 00000000..7756fde2 --- /dev/null +++ b/backends/cpu/compiler.cmake @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +target_sources(_compiler PRIVATE cpu/compiler_cpu.cpp cpu/ApproximateMath.cpp) +target_include_directories(_compiler SYSTEM PRIVATE + "${LLVM_MAIN_SRC_DIR}/../lld/include") +target_link_libraries(_compiler PRIVATE + MLIRMemRefTransforms MLIRMathTransforms + MLIRAffineToStandard MLIRVectorToSCF + MLIRSCFToControlFlow MLIRVectorToLLVMPass + MLIRMathToLibm MLIRConvertToLLVMPass + MLIRReconcileUnrealizedCasts + MLIRTargetLLVMIRExport lldELF lldCommon ${CMAKE_DL_LIBS} m) +llvm_map_components_to_libnames(_llvm_libs nativecodegen passes) +target_link_libraries(_compiler PRIVATE ${_llvm_libs}) + +if(NOT LLVM_ENABLE_RTTI) + set_source_files_properties(cpu/ApproximateMath.cpp + PROPERTIES COMPILE_OPTIONS "-fno-rtti") +endif() \ No newline at end of file diff --git a/backends/cpu/compiler_cpu.cpp b/backends/cpu/compiler_cpu.cpp new file mode 100644 index 00000000..fc943caf --- /dev/null +++ b/backends/cpu/compiler_cpu.cpp @@ -0,0 +1,332 @@ +// SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights +// reserved. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "ApproximateMath.h" +#include "amx.h" +#include "lld/Common/Driver.h" +#include "mlir/Conversion/Passes.h" +#include "mlir/Conversion/TileIRToMLIR/Passes.h" +#include "mlir/Conversion/VectorToSCF/VectorToSCF.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Math/Transforms/Passes.h" +#include "mlir/Dialect/MemRef/Transforms/Passes.h" +#include "mlir/Dialect/Vector/Transforms/Passes.h" +#include "mlir/LowerVectorScan.h" +#include "mlir/Pass/PassRegistry.h" +#include "mlir/Target/LLVMIR/Export.h" +#include "mlir/Transforms/Passes.h" +#include "mlir/compiler.h" +#include "llvm/IR/LegacyPassManager.h" +#include "llvm/IR/Module.h" +#include "llvm/IR/Verifier.h" +#include "llvm/MC/TargetRegistry.h" +#include "llvm/Passes/PassBuilder.h" +#include "llvm/Support/CodeGen.h" +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/MemoryBuffer.h" +#include "llvm/Support/Path.h" +#include "llvm/Support/TargetSelect.h" +#include "llvm/Target/TargetMachine.h" +#include "llvm/TargetParser/Host.h" +#include "llvm/TargetParser/SubtargetFeature.h" +#include "llvm/TargetParser/Triple.h" +#include +#include +#include +#include + +LLD_HAS_DRIVER(elf) + +namespace cutile { + +const char *const backendName = "cpu"; + +void registerBackendTranslations(mlir::DialectRegistry &) {} + +void registerBackendPasses() { + mlir::registerTransformsPasses(); + mlir::registerLowerAffinePass(); + mlir::registerFinalizeMemRefToLLVMConversionPass(); + mlir::registerSCFToControlFlowPass(); + mlir::registerConvertIndexToLLVMPass(); + mlir::registerConvertVectorToLLVMPass(); + mlir::registerConvertVectorToSCF(); + mlir::registerConvertMathToLLVMPass(); + mlir::registerConvertMathToLibmPass(); + mlir::registerArithToLLVMConversionPass(); + mlir::registerUBToLLVMConversionPass(); + mlir::registerConvertToLLVMPass(); + mlir::registerReconcileUnrealizedCastsPass(); + mlir::math::registerMathPasses(); + mlir::memref::registerMemRefPasses(); + mlir::vector::registerVectorPasses(); + if (llvm::InitializeNativeTarget() || + llvm::InitializeNativeTargetAsmPrinter()) + throw std::runtime_error( + "cuTile compiler: native LLVM target is unavailable"); +} + +static llvm::StringRef cpuTransformEntryPoint() { + static const char *entryPoint = [] { + if (llvm::Triple(llvm::sys::getProcessTriple()).getArch() != + llvm::Triple::x86_64) + return "__transform_nano"; + return cutile::cpu::enableAMX() ? "__transform_nano_amx" + : "__transform_nano_x86"; + }(); + return entryPoint; +} + +static std::unique_ptr +lowerToLLVM(Compilation &compilation, mlir::ModuleOp module, + llvm::StringRef entryPoint, llvm::LLVMContext &llvmContext) { + mlir::PassManager pm(&compilation.context); + + entryPoint = entryPoint.empty() ? cpuTransformEntryPoint() : entryPoint; + if (entryPoint != "__transform_nano" && + entryPoint != "__transform_nano_x86" && + entryPoint != "__transform_nano_amx") + throw std::invalid_argument("unknown built-in CPU pipeline entry point"); + pm.addPass(createPrepareCPUVectorsPass(entryPoint)); + pm.addPass(createLowerVectorScanPass()); + // Must run before vector lowering: LLVM scalarizes vector math intrinsics + // into per-element libm calls. + pm.addPass(createApproximateMathPass()); + + pm.addPass(mlir::createCanonicalizerPass()); + pm.addPass(mlir::createLoopInvariantCodeMotionPass()); + pm.addPass(mlir::createCSEPass()); + pm.addPass(mlir::createSymbolDCEPass()); + pm.nest().addPass( + mlir::vector::createLowerVectorMultiReductionPass( + mlir::vector::VectorMultiReductionLowering::InnerReduction)); + pm.addPass(mlir::memref::createFoldMemRefAliasOpsPass()); + pm.addPass(mlir::memref::createExpandStridedMetadataPass()); + pm.addPass(mlir::createCanonicalizerPass()); + + mlir::VectorTransferToSCFOptions vectorToSCFOptions; + vectorToSCFOptions.enableFullUnroll().setTargetRank(1); + pm.addPass(mlir::createConvertVectorToSCFPass(vectorToSCFOptions)); + pm.addPass(mlir::createLowerAffinePass()); + pm.addPass(mlir::createFinalizeMemRefToLLVMConversionPass()); + pm.addPass(mlir::createSCFToControlFlowPass()); + pm.addPass(mlir::createConvertIndexToLLVMPass()); + + mlir::ConvertVectorToLLVMPassOptions vectorToLLVMOptions; + vectorToLLVMOptions.reassociateFPReductions = true; + vectorToLLVMOptions.vectorContractLowering = + mlir::vector::VectorContractLowering::OuterProduct; + pm.addPass( + mlir::createConvertVectorToLLVMPass(std::move(vectorToLLVMOptions))); + + mlir::math::MathExpandOpsPassOptions mathExpandOptions; + mathExpandOptions.opMnemonics.push_back("fpowi"); + pm.addPass(mlir::math::createMathExpandOpsPass(std::move(mathExpandOptions))); + mlir::ConvertMathToLLVMPassOptions mathToLLVMOptions; + mathToLLVMOptions.approximateLog1p = true; + pm.addPass(mlir::createConvertMathToLLVMPass(std::move(mathToLLVMOptions))); + pm.addPass(mlir::createConvertMathToLibmPass()); + pm.addPass(mlir::createReconcileUnrealizedCastsPass()); + pm.addPass(mlir::createArithToLLVMConversionPass()); + pm.addPass(mlir::createUBToLLVMConversionPass()); + pm.addPass(mlir::createConvertToLLVMPass()); + pm.addPass(mlir::createReconcileUnrealizedCastsPass()); + pm.addPass(mlir::createCanonicalizerPass()); + pm.addPass(mlir::createCSEPass()); + pm.addPass(mlir::createSymbolDCEPass()); + + compilation.run(module, pm); + auto result = mlir::translateModuleToLLVMIR(module, llvmContext); + if (!result) + compilation.fail("translating to LLVM IR"); + return result; +} + +static void writeFile(llvm::StringRef path, llvm::StringRef data) { + std::error_code error; + llvm::raw_fd_ostream output(path, error); + if (error) + throw std::runtime_error("cuTile compiler: " + path.str() + ": " + + error.message()); + output << data; + output.close(); + if (output.has_error()) + throw std::runtime_error("cuTile compiler: unable to write " + path.str()); +} + +// Resolve by SONAME: a symbol address would name an interposer such as an +// LD_PRELOADed allocator instead of the system library. +static std::string loadedLibrary(const char *soname) { + void *handle = dlopen(soname, RTLD_LAZY | RTLD_NOLOAD); + link_map *map = nullptr; + std::string path; + if (handle && dlinfo(handle, RTLD_DI_LINKMAP, &map) == 0 && map) + path = map->l_name; + if (handle) + dlclose(handle); + if (path.empty()) + throw std::runtime_error(std::string("cuTile compiler: ") + soname + + " is not loaded"); + return path; +} + +static std::string linkObject(llvm::StringRef object, + llvm::StringRef directory) { + llvm::SmallString<128> objectPath(directory), sharedPath(directory); + llvm::sys::path::append(objectPath, "kernel.o"); + llvm::sys::path::append(sharedPath, "kernel.so"); + writeFile(objectPath, object); + std::string math = loadedLibrary("libm.so.6"); + std::string libc = loadedLibrary("libc.so.6"); + std::string runtime = loadedLibrary("libgcc_s.so.1"); + // One thread: LLD's default thread pool would persist in the host process + // and deadlock links in fork()ed children. + const char *args[] = {"ld.lld", + "-shared", + "-z", + "defs", + "--threads=1", + "--as-needed", + objectPath.c_str(), + math.c_str(), + libc.c_str(), + runtime.c_str(), + "-o", + sharedPath.c_str()}; + // LLD has process-global state. It recovers from its fatal errors, but must + // not run again in the same process afterwards. + static std::mutex linkerMutex; + static bool linkerUsable = true; + std::lock_guard guard(linkerMutex); + if (!linkerUsable) + throw std::runtime_error( + "cuTile compiler: the linker failed fatally before; restart the " + "process or set a compiler timeout to link in a worker process"); + std::string output, errors; + llvm::raw_string_ostream stdoutStream(output), stderrStream(errors); + lld::DriverDef drivers[] = {{lld::Gnu, &lld::elf::link}}; + lld::Result result = lld::lldMain(args, stdoutStream, stderrStream, drivers); + linkerUsable = result.canRunAgain; + if (result.retCode) + throw std::runtime_error("cuTile compiler: linking failed\n" + errors + + output); + if (!errors.empty()) + llvm::errs() << errors; + auto binary = llvm::MemoryBuffer::getFile(sharedPath); + if (!binary) + throw std::runtime_error("cuTile compiler: unable to read linked kernel: " + + binary.getError().message()); + return (*binary)->getBuffer().str(); +} + +static std::string compileCPU(llvm::StringRef bytecode, bool assumeInBounds, + llvm::StringRef directory, + llvm::StringRef dumpPath, bool printIR, + bool packArguments) { + Compilation compilation(printIR); + auto module = compilation.readTileIR(bytecode); + compilation.lowerTileIR( + *module, + "target=cpu append-grid-args=true " + "assume-in-bounds=" + + std::string(assumeInBounds ? "true" : "false"), + [&](mlir::OpPassManager &pm) { + pm.addPass(createLowerVectorScanPass()); + pm.addPass(mlir::createLoopInvariantCodeMotionPass()); + pm.addPass(mlir::createConvertMemrefArgsToPtrArgsPass()); + if (packArguments) + pm.addPass(mlir::createPackCPUKernelArgumentsPass()); + pm.addPass(mlir::createCSEPass()); + pm.addPass(mlir::createCanonicalizerPass()); + }); + if (!dumpPath.empty()) + writeFile(dumpPath, printModule(*module)); + llvm::LLVMContext llvmContext; + auto llvmModule = lowerToLLVM(compilation, *module, {}, llvmContext); + auto triple = llvm::sys::getProcessTriple(); + std::string error; + auto *target = + llvm::TargetRegistry::lookupTarget(llvm::Triple(triple), error); + if (!target) + throw std::runtime_error("cuTile compiler: " + error); + llvm::SubtargetFeatures features; + for (const auto &feature : llvm::sys::getHostCPUFeatures()) + features.AddFeature(feature.first(), feature.second); + std::unique_ptr machine(target->createTargetMachine( + llvm::Triple(triple), llvm::sys::getHostCPUName(), features.getString(), + llvm::TargetOptions(), llvm::Reloc::PIC_, std::nullopt, + llvm::CodeGenOptLevel::Aggressive)); + if (!machine) + throw std::runtime_error("cuTile compiler: unable to create native target"); + llvmModule->setTargetTriple(llvm::Triple(triple)); + llvmModule->setDataLayout(machine->createDataLayout()); + llvm::PassBuilder builder(llvmContext, machine.get()); + llvm::LoopAnalysisManager loops; + llvm::FunctionAnalysisManager functions; + llvm::CGSCCAnalysisManager callGraph; + llvm::ModuleAnalysisManager modules; + builder.registerLoopAnalyses(loops); + builder.registerFunctionAnalyses(functions); + builder.registerCGSCCAnalyses(callGraph); + builder.registerModuleAnalyses(modules); + builder.crossRegisterProxies(loops, functions, callGraph, modules); + auto optimizer = + builder.buildPerModuleDefaultPipeline(llvm::OptimizationLevel::O3); + optimizer.run(*llvmModule, modules); + llvm::raw_string_ostream errors(compilation.diagnostics); + if (llvm::verifyModule(*llvmModule, &errors)) + compilation.fail("verifying LLVM IR"); + llvm::SmallVector object; + llvm::raw_svector_ostream stream(object); + llvm::legacy::PassManager codegen; + if (machine->addPassesToEmitFile(codegen, stream, nullptr, + llvm::CodeGenFileType::ObjectFile)) + throw std::runtime_error( + "cuTile compiler: native object emission is unavailable"); + codegen.run(*llvmModule); + return linkObject(llvm::StringRef(object.data(), object.size()), directory); +} + +void bindBackend(nanobind::module_ &module) { + namespace nb = nanobind; + using namespace nb::literals; + module.def( + "compile_cpu", + [](nb::bytes bytecode, bool assumeInBounds, const std::string &directory, + const std::string &dumpPath, bool printIR, bool packArguments) { + llvm::StringRef input(bytecode.c_str(), bytecode.size()); + std::string binary; + { + nb::gil_scoped_release release; + binary = compileCPU(input, assumeInBounds, directory, dumpPath, + printIR, packArguments); + } + return nb::bytes(binary.data(), binary.size()); + }, + "bytecode"_a, "assume_in_bounds"_a, "directory"_a, "dump_path"_a = "", + "print_ir"_a = false, "pack_arguments"_a = true); + module.def( + "lower_cpu_mlir", + [](nb::bytes source, const std::string &entryPoint) { + llvm::StringRef input(source.c_str(), source.size()); + std::string result; + { + nb::gil_scoped_release release; + Compilation compilation; + auto module = compilation.readMLIR(input); + llvm::LLVMContext llvmContext; + auto llvmModule = + lowerToLLVM(compilation, *module, entryPoint, llvmContext); + llvm::raw_string_ostream stream(result); + llvmModule->print(stream, nullptr); + } + return nb::bytes(result.data(), result.size()); + }, + "source"_a, "entry_point"_a = ""); + module.def("cpu_transform_entry_point", + [] { return cpuTransformEntryPoint().str(); }); +} + +} // namespace cutile diff --git a/backends/cpu/launcher_cpu.cpp b/backends/cpu/launcher_cpu.cpp new file mode 100644 index 00000000..fbcdb4bb --- /dev/null +++ b/backends/cpu/launcher_cpu.cpp @@ -0,0 +1,142 @@ +// SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights +// reserved. +// +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include + +#include "amx.h" +#include "../launch.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace nb = nanobind; +using namespace nb::literals; + +namespace { + +struct LibraryCloser { + void operator()(void *handle) const { dlclose(handle); } +}; + +class Kernel { +public: + static constexpr bool memref = false; + static constexpr int device = nb::device::cpu::value; + using Arguments = std::vector; + + Kernel(const std::string &path, const std::string &symbol, + size_t argumentBytes) + : library(dlopen(path.c_str(), RTLD_NOW | RTLD_LOCAL)), + argumentBytes(argumentBytes) { + if (!library) + throw std::runtime_error( + std::string("CPU backend: unable to load kernel (TMPDIR must allow " + "executable files): ") + + dlerror()); + dlerror(); + function = reinterpret_cast(dlsym(library.get(), symbol.c_str())); + const char *error = dlerror(); + if (error || !function) + throw std::runtime_error("CPU backend: unable to find kernel '" + symbol + + "': " + (error ? error : "null entry point")); + if (argumentBytes % 8 != 0) + throw nb::value_error("CPU argument buffer size must be a multiple of 8"); + } + + void launch(nb::bytes arguments, const std::array &grid, + int64_t numThreads) const { + if (arguments.size() != argumentBytes) + throw nb::value_error( + "CPU argument buffer size does not match the kernel"); + run(arguments.data(), grid, numThreads); + } + + void launchNative(const Arguments &arguments, + const std::array &grid, nb::handle options, + nb::handle) const { + if (arguments.size() * 8 != argumentBytes) + throw nb::value_error("CPU argument buffer size does not match the kernel"); + nb::object setting = nb::steal( + PyMapping_GetItemString(options.ptr(), "num_cpu_threads")); + if (!setting.is_valid() && !PyErr_ExceptionMatches(PyExc_KeyError)) + throw nb::python_error(); + if (!setting.is_valid()) + PyErr_Clear(); + int64_t threads = setting.is_valid() ? cutile::launch::index(setting) : 0; + run(arguments.data(), grid, threads); + } + +private: + void run(const void *data, const std::array &grid, + int64_t numThreads) const { + for (int64_t dimension : grid) { + if (dimension < 0 || dimension > std::numeric_limits::max()) + throw nb::value_error( + "CPU grid dimensions must be nonnegative int32 values"); + } + if (numThreads < 0 || numThreads > std::numeric_limits::max()) + throw nb::value_error( + "num_cpu_threads must be a nonnegative int32 value"); + if (grid[0] == 0 || grid[1] == 0 || grid[2] == 0) + return; + int64_t blocks = 1; + for (int64_t dimension : grid) { + if (blocks > std::numeric_limits::max() / dimension) + throw nb::value_error("CPU grid block count must fit in int64"); + blocks *= dimension; + } + + const int32_t gridX = grid[0], gridY = grid[1], gridZ = grid[2]; + nb::gil_scoped_release release; + if (blocks == 1) { + function(data, 0, 0, 0, 1, 1, 1); + return; + } + const int threads = numThreads > 0 ? numThreads : omp_get_max_threads(); + auto invoke = [&](int64_t index) { + const int32_t x = index % gridX; + const int32_t y = (index / gridX) % gridY; + const int32_t z = index / (int64_t(gridX) * gridY); + function(data, x, y, z, gridX, gridY, gridZ); + }; + if (threads == 1) { + for (int64_t index = 0; index < blocks; ++index) + invoke(index); + return; + } +#pragma omp parallel for schedule(static) num_threads(threads) + for (int64_t index = 0; index < blocks; ++index) + invoke(index); + } + +private: + // The compiler loads each argument from an eight-byte slot. Grid coordinates + // stay separate so all blocks can share the immutable argument buffer. + using Entry = void (*)(const void *, int32_t, int32_t, int32_t, int32_t, + int32_t, int32_t); + std::unique_ptr library; + Entry function; + size_t argumentBytes; +}; + +} // namespace + +NB_MODULE(_cpu, module) { + cutile::launch::bindCache(module); + module.attr("amx_enabled") = cutile::cpu::enableAMX(); + nb::class_(module, "Kernel") + .def(nb::init(), + "path"_a, "symbol"_a, "argument_bytes"_a) + .def("launch", &Kernel::launch, "arguments"_a, "grid"_a, + "num_cpu_threads"_a = 0); +} diff --git a/backends/cpu/test/bench_launch.py b/backends/cpu/test/bench_launch.py new file mode 100644 index 00000000..81592b88 --- /dev/null +++ b/backends/cpu/test/bench_launch.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import statistics +import timeit + +import torch + +import cuda.tile as ct +from cuda.tile._backends import cpu +from cuda.tile._cext import default_tile_context + + +@ct.kernel +def vector_add(a, b, out, block: ct.Constant[int]): + index = ct.bid(0) + left = ct.load(a, index=(index,), shape=(block,)) + right = ct.load(b, index=(index,), shape=(block,)) + ct.store(out, index=(index,), tile=left + right) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--disable-native", action="store_true") + parser.add_argument("--compare", action="store_true") + parser.add_argument("--iterations", type=int, default=10000) + args = parser.parse_args() + default_tile_context.config.cache_dir = None + if args.disable_native: + cpu.make_launch_cache = None + factory = getattr(cpu, "make_launch_cache", None) + torch.set_num_threads(1) + ct.set_backend("cpu") + print(f"native cache: {bool(getattr(cpu, 'make_launch_cache', None))}") + for blocks, threads in ((0, 1), (1, 1), (64, 1), (64, 4)): + size = max(blocks, 1) * 256 + left, right = torch.randn(size), torch.randn(size) + out = torch.empty_like(left) + values = (left, right, out, 256) + grid = (blocks,) + with ct.backend_options({"num_cpu_threads": threads}): + for _ in range(20): + ct.launch(None, grid, vector_add, values) + samples = {"python": [], "native": []} if args.compare else {"selected": []} + for repetition in range(9): + modes = list(samples) + if repetition % 2: + modes.reverse() + for mode in modes: + cpu.make_launch_cache = None if mode == "python" else factory + ct.launch(None, grid, vector_add, values) + elapsed = timeit.timeit( + lambda: ct.launch(None, grid, vector_add, values), + number=args.iterations) + samples[mode].append(elapsed * 1e6 / args.iterations) + if blocks: + torch.testing.assert_close(out, left + right) + medians = {mode: statistics.median(times) for mode, times in samples.items()} + result = " ".join(f"{mode}={median:.3f} us" for mode, median in medians.items()) + if args.compare: + result += f" speedup={medians['python'] / medians['native']:.2f}x" + print(f"blocks={blocks:2} threads={threads}: {result}") + cpu.make_launch_cache = factory + + +if __name__ == "__main__": + main() diff --git a/backends/cpu/test/test_compiler_cpu.py b/backends/cpu/test/test_compiler_cpu.py new file mode 100644 index 00000000..41cbad17 --- /dev/null +++ b/backends/cpu/test/test_compiler_cpu.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from concurrent.futures import ThreadPoolExecutor +import os +import shutil +import signal +import subprocess +import sys +import warnings + +import numpy as np +import pytest + +import cuda.tile as ct +from cuda.tile._backends import _custom, _native, cpu +from cuda.tile._backends._signature import constant_parameters +from cuda.tile._cext import default_tile_context + + +@pytest.fixture(autouse=True) +def _cpu_compiler(native): + if native.backend != "cpu": + pytest.skip("requires the CPU compiler") + + +@ct.kernel +def _waves(x, out, TILE: ct.Constant[int]): + value = ct.load(x, index=(ct.bid(0),), shape=(TILE,)) + ct.store(out, index=(ct.bid(0),), tile=ct.sin(value) + ct.cos(value)) + + +@ct.kernel +def _empty(): + pass + + +@pytest.mark.parametrize("concurrent", [False, True]) +def test_cpu_compiles_and_links_without_external_tools(native, bytecode, monkeypatch, concurrent): + monkeypatch.setenv("PATH", "") + monkeypatch.setattr( + _native.subprocess, "run", + lambda *args, **kwargs: pytest.fail("unexpected compiler subprocess")) + x = np.linspace(-1, 1, 128, dtype=np.float32) + out = np.empty_like(x) + arguments = (x, out, 16) + data, signature = bytecode(_waves, arguments, cpu) + options = cpu.Options(False) + identity = cpu.compiler_identity(options) + assert isinstance(identity, str) and len(identity) == 64 + if concurrent: + with ThreadPoolExecutor(max_workers=4) as pool: + binaries = list(pool.map(lambda _: cpu.compile(data, signature, options), range(4))) + else: + binaries = [cpu.compile(data, signature, options) for _ in range(3)] + constants = constant_parameters(_waves) + runtime = _custom._runtime_arguments( + arguments, constants, _custom._describe(constants, arguments)) + for binary in binaries: + assert binary.startswith(b"\x7fELF") + out.fill(np.nan) + loaded = cpu.load(binary, signature, options) + cpu.launch(loaded, None, (8,), runtime) + np.testing.assert_allclose(out, np.sin(x) + np.cos(x), rtol=2e-6, atol=2e-6) + + +def test_invalid_bytecode_does_not_poison_compiler(native, bytecode, tmp_path): + with pytest.raises(ValueError, match="invalid TileIR"): + native.compile_cpu(b"invalid", False, str(tmp_path)) + data, signature = bytecode(_empty, (), cpu) + assert cpu.compile(data, signature, cpu.Options(False)).startswith(b"\x7fELF") + + +def test_native_mlir_errors_include_diagnostics(native): + with pytest.raises(RuntimeError, match="parsing MLIR"): + native.lower_cpu_mlir(b"not MLIR") + + +@pytest.mark.parametrize("entry_point", [ + "__transform_nano", "__transform_nano_x86", "__transform_nano_amx"]) +def test_builtin_cpu_pipeline_lowers_vectors(native, entry_point): + source = b""" +module { + func.func @identity(%value: vector<1x4xf32>) -> vector<1x4xf32> { + %result = arith.addf %value, %value : vector<1x4xf32> + return %result : vector<1x4xf32> + } +} +""" + direct = native.lower_cpu_mlir(source, entry_point) + assert b"@identity" in direct + + +def test_builtin_cpu_pipeline_rejects_unknown_entry_point(native): + with pytest.raises(ValueError, match="unknown built-in CPU pipeline entry point"): + native.lower_cpu_mlir(b"module {}", "unknown") + + +@ct.kernel +def _divide(x, y, out): + a = ct.load(x, index=(0,), shape=(1,)) + b = ct.load(y, index=(0,), shape=(1,)) + ct.store(out, index=(0,), tile=ct.truediv(a, b, rounding_mode=ct.RoundingMode.RN)) + + +def test_native_cpu_default_rounding_still_executes(native, bytecode): + x, y = np.array([1.0], dtype=np.float32), np.array([3.0], dtype=np.float32) + out = np.zeros_like(x) + data, signature = bytecode(_divide, (x, y, out), cpu) + loaded = cpu.load(cpu.compile(data, signature, cpu.Options(False)), + signature, cpu.Options(False)) + runtime = _custom._runtime_arguments( + (x, y, out), (False,) * 3, _custom._describe((False,) * 3, (x, y, out))) + cpu.launch(loaded, None, (1,), runtime) + np.testing.assert_array_equal(out, x / y) + + +def test_cpu_compiles_in_forked_child_after_parent_compiled(native, bytecode): + x = np.zeros(128, dtype=np.float32) + data, signature = bytecode(_waves, (x, np.empty_like(x), 16), cpu) + options = cpu.Options(False) + cpu.compile(data, signature, options) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + pid = os.fork() + if pid == 0: + status = 1 + try: + signal.alarm(60) + status = 0 if cpu.compile(data, signature, options).startswith(b"\x7fELF") else 2 + finally: + os._exit(status) + _, status = os.waitpid(pid, 0) + assert os.waitstatus_to_exitcode(status) == 0 + + +_MALLOC_SHIM = b""" +#define _GNU_SOURCE +#include +#include +void *malloc(size_t size) { + static void *(*next)(size_t); + if (!next) + next = (void *(*)(size_t))dlsym(RTLD_NEXT, "malloc"); + return next(size); +} +""" + +_PRELOADED_COMPILE = """\ +import sys, tempfile +from cuda.tile._backends import _compiler +with tempfile.TemporaryDirectory() as directory: + sys.stdout.buffer.write(_compiler.compile_cpu( + sys.stdin.buffer.read(), False, directory)) +""" + + +def test_cpu_links_system_libraries_despite_preloaded_allocator(native, bytecode, tmp_path): + compiler = shutil.which("cc") + if compiler is None: + pytest.skip("building the malloc shim requires cc") + shim = tmp_path / "libmalloc-shim.so" + subprocess.run([compiler, "-shared", "-fPIC", "-x", "c", "-", "-o", str(shim), "-ldl"], + input=_MALLOC_SHIM, check=True) + x = np.zeros(128, dtype=np.float32) + data, _ = bytecode(_waves, (x, np.empty_like(x), 16), cpu) + binary = subprocess.run( + [sys.executable, "-c", _PRELOADED_COMPILE], input=data, + env={**os.environ, "LD_PRELOAD": str(shim)}, + capture_output=True, check=True, timeout=300).stdout + assert binary.startswith(b"\x7fELF") + assert b"libm.so.6" in binary + assert str(shim).encode() not in binary + + +@pytest.mark.parametrize("timeout", [None, 30]) +def test_native_errors_with_and_without_worker(native, monkeypatch, tmp_path, timeout): + monkeypatch.setattr(default_tile_context.config, "compiler_timeout_sec", timeout) + with pytest.raises(ValueError, match="invalid TileIR bytecode") as caught: + _native.invoke("cpu", "compile_cpu", b"invalid", False, str(tmp_path)) + assert "Traceback" not in str(caught.value) + with pytest.raises(RuntimeError, match="parsing MLIR") as caught: + _native.invoke("cpu", "lower_cpu_mlir", b"not MLIR") + assert "Traceback" not in str(caught.value) diff --git a/backends/cpu/test/test_cpu_backend.py b/backends/cpu/test/test_cpu_backend.py new file mode 100644 index 00000000..6459536a --- /dev/null +++ b/backends/cpu/test/test_cpu_backend.py @@ -0,0 +1,960 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import ctypes +from concurrent.futures import ThreadPoolExecutor +import importlib.util +import os +from pathlib import Path +import shutil +import struct +import subprocess +import sys +from types import SimpleNamespace +from typing import Annotated + +import cuda.tile as ct +from cuda.tile import _backends as _backend +from cuda.tile._backends import _custom, cpu +from cuda.tile._backends._signature import ( + build_signature, constant_parameters, runtime_parameters, +) +from cuda.tile._exception import TypeCheckingError +import numpy as np +import pytest +import torch + +ConstInt = ct.Constant[int] + + +@ct.kernel +def _vector_add(a, b, c, tile_size: ConstInt): + block = ct.bid(0) + a_tile = ct.load(a, index=(block,), shape=(tile_size,)) + b_tile = ct.load(b, index=(block,), shape=(tile_size,)) + ct.store(c, index=(block,), tile=a_tile + b_tile) + + +def _arguments(): + return ( + np.empty(64, dtype=np.float32), + np.empty(64, dtype=np.float32), + np.empty(64, dtype=np.float32), + 16, + ) + + +def _runtime(kernel, args): + """The runtime arguments ``ct.launch`` passes to the backend.""" + constants = constant_parameters(kernel) + return _custom._runtime_arguments( + args, constants, _custom._describe(constants, args)) + + +def _parameters(kernel, args): + return runtime_parameters("CPU", build_signature(kernel, args)) + + +def test_signature_flattens_arrays_and_omits_constants(): + arguments = _arguments() + parameters = _parameters(_vector_add, arguments) + + values = cpu._flatten_arguments(parameters, _runtime(_vector_add, arguments)) + + assert cpu._argument_formats(parameters) == ["P", "i", "i"] * 3 + assert values[0] == arguments[0].ctypes.data + assert values[1:3] == [64, 1] + assert values[4:6] == [64, 1] + assert values[7:9] == [64, 1] + + +def test_signature_uses_actual_array_alignment(): + arguments = list(_arguments()) + arguments[0] = arguments[0][1:] + + signature = build_signature(_vector_add, tuple(arguments)) + + assert signature.parameters[0].base_addr_divisible_by == 1 + + +def test_signature_rejects_negative_strides(): + arguments = list(_arguments()) + arguments[0] = arguments[0][::-1] + + with pytest.raises(ValueError, match="negative array strides"): + build_signature(_vector_add, tuple(arguments)) + + +@ct.kernel +def _scalar_and_array(x, flag: bool, scale: float, count: int): + pass + + +def test_signature_matches_native_launcher_types(): + x = np.zeros(4, dtype=np.bool_) + + parameters = build_signature(_scalar_and_array, (x, True, 1.0, 1)).parameters + + assert parameters[0].dtype is ct.bool_ + assert [p.dtype for p in parameters[1:]] == [ct.bool_, ct.float32, ct.int32] + + +def test_signature_detects_internally_aliasing_arrays(): + base = np.zeros(16, dtype=np.float32) + broadcast = np.lib.stride_tricks.as_strided(base, shape=(4, 4), strides=(0, 4)) + + aliasing = build_signature(_scalar_and_array, (broadcast, True, 1.0, 1)) + disjoint = build_signature(_scalar_and_array, (base.reshape(4, 4), True, 1.0, 1)) + + assert aliasing.parameters[0].may_alias_internally + assert not disjoint.parameters[0].may_alias_internally + + +def test_signature_skips_layout_assumptions_beyond_five_dimensions(): + x = np.zeros((16,) * 6, dtype=np.float32) + + [array, *_] = build_signature(_scalar_and_array, (x, True, 1.0, 1)).parameters + + assert array.stride_constant == (None,) * 6 + assert array.shape_divisible_by == (1,) * 6 + + +def test_flatten_arguments_rejects_non_host_arrays(): + arguments = _arguments() + device_tensor = torch.empty(64, dtype=torch.float32, device="meta") + + with pytest.raises(ValueError, match="expected 'cpu'"): + cpu._flatten_arguments(_parameters(_vector_add, arguments), + _runtime(_vector_add, (device_tensor, *arguments[1:]))) + + +def test_flatten_arguments_rejects_mismatched_kinds(): + arguments = _arguments() + runtime = _runtime(_vector_add, arguments) + + with pytest.raises(TypeError, match="expected a 1-D array"): + cpu._flatten_arguments(_parameters(_vector_add, arguments), (1, *runtime[1:])) + with pytest.raises(ValueError): + cpu._flatten_arguments(_parameters(_vector_add, arguments), runtime[:2]) + + +def test_signature_rejects_array_lists(): + with pytest.raises(TypeError, match="only supported on CUDA"): + build_signature(_scalar_and_array, ([np.zeros(4)], True, 1.0, 1)) + + +@pytest.mark.usefixtures("native_cpu") +def test_native_compiler_selects_builtin_entry_point(): + if cpu.platform.machine() != "x86_64": + assert cpu._transform_entry_point() == "__transform_nano" + elif cpu._amx_enabled(): + assert cpu._transform_entry_point() == "__transform_nano_amx" + else: + assert cpu._transform_entry_point() == "__transform_nano_x86" + + +@pytest.mark.usefixtures("native_cpu") +def test_native_lowering_pipeline_builds_pass_manager(): + native = cpu.native_compiler("cpu") + source = b""" +module { + func.func @identity(%value: i32) -> i32 { + %result = arith.addi %value, %value : i32 + return %result : i32 + } +} +""" + llvm_ir = native.lower_cpu_mlir( + source, "__transform_nano") + assert b"define i32 @identity" in llvm_ir + + +@pytest.mark.usefixtures("native_cpu") +def test_native_lowering_pipeline_lowers_vector_scan(): + native = cpu.native_compiler("cpu") + source = b""" +module { + func.func @scan(%input: vector<4xf32>) -> vector<4xf32> { + %zero = arith.constant dense<0.0> : vector + %scan, %acc = vector.scan , %input, %zero reduction_dim = 0, inclusive = true : + vector<4xf32>, vector + return %scan : vector<4xf32> + } +} +""" + llvm_ir = native.lower_cpu_mlir( + source, "__transform_nano") + assert b"@scan" in llvm_ir + assert b"vector.scan" not in llvm_ir + + +def test_backend_registration_accepts_builtin_cpu(): + _backend.set_backend("cpu") + try: + assert _backend.get_backend() is cpu + finally: + _backend.set_backend(None) + + +def test_compiler_identity_tracks_toolchain_cpu_and_compile_options(monkeypatch): + monkeypatch.setattr( + cpu, "native_compiler", lambda _: SimpleNamespace(__file__="/compiler")) + monkeypatch.setattr( + cpu, "file_fingerprint", lambda path: f"fingerprint:{path}") + monkeypatch.setattr(cpu, "_host_cpu", lambda: {"flags": "avx2"}) + monkeypatch.setattr(cpu, "_transform_entry_point", lambda: "__transform_nano_x86") + + first = cpu.compiler_identity(cpu.normalize_options({"num_cpu_threads": 1})) + second = cpu.compiler_identity(cpu.normalize_options({"num_cpu_threads": 8})) + in_bounds = cpu.compiler_identity( + cpu.normalize_options({"assume_in_bounds": True})) + + assert first is not None + assert first == second + assert first != in_bounds + + monkeypatch.setattr(cpu, "_host_cpu", lambda: {"flags": "avx2 avx512f"}) + assert first != cpu.compiler_identity(cpu.Options(False)) + monkeypatch.setattr(cpu, "_host_cpu", lambda: {"flags": "avx2"}) + # Whether AMX is enabled selects the pipeline, so it is part of the identity. + monkeypatch.setattr(cpu, "_transform_entry_point", lambda: "__transform_nano_amx") + assert first != cpu.compiler_identity(cpu.Options(False)) + monkeypatch.setattr(cpu, "_transform_entry_point", lambda: "__transform_nano_x86") + monkeypatch.setattr( + cpu, "file_fingerprint", lambda path: f"rebuilt:{path}") + assert first != cpu.compiler_identity(cpu.Options(False)) + + +def test_compiler_identity_disables_disk_cache_for_missing_native_file(monkeypatch): + monkeypatch.setattr( + cpu, "native_compiler", lambda _: SimpleNamespace(__file__="/compiler")) + + def missing_file(*_args): + raise FileNotFoundError("missing compiler") + + monkeypatch.setattr(cpu, "file_fingerprint", missing_file) + assert cpu.compiler_identity(cpu.Options(False)) is None + + +def test_compile_calls_native_pipeline_and_cleans_temporary_files(monkeypatch): + observed = {} + monkeypatch.setattr(cpu, "_transform_entry_point", lambda: "__transform_nano") + monkeypatch.setenv("CUTILE_CPU_DUMP_MLIR", "/dump.mlir") + monkeypatch.setenv("MLIR_ENABLE_DUMP", "1") + + def invoke(backend, method, bytecode, in_bounds, directory, dump, + print_ir): + observed.update(backend=backend, method=method, bytecode=bytecode, + in_bounds=in_bounds, + directory=directory, dump=dump, print_ir=print_ir) + assert Path(directory).is_dir() + return b"shared object" + + monkeypatch.setattr(cpu, "invoke", invoke) + signature = build_signature(_vector_add, _arguments()) + assert cpu.compile(b"bytecode", signature, cpu.Options(True)) == b"shared object" + assert observed["backend"] == "cpu" + assert observed["method"] == "compile_cpu" + assert observed["bytecode"] == b"bytecode" + assert observed["in_bounds"] and observed["print_ir"] + assert observed["dump"] == "/dump.mlir" + assert not Path(observed["directory"]).exists() + + +def test_options_reject_non_bool_flags(): + with pytest.raises(TypeError, match="must be a bool"): + cpu.normalize_options({"assume_in_bounds": "false"}) + + +def test_options_reject_unknown_keys(monkeypatch): + with pytest.raises(ValueError, match="unknown options 'assume_inbounds'"): + cpu.normalize_options({"assume_inbounds": True}) + monkeypatch.setenv("CUTILE_CPU_ASSUME_IN_BOUNDS", "1") + assert cpu.normalize_options({}) == cpu.Options(assume_in_bounds=False) + + +def test_load_prepares_native_launcher_once(monkeypatch): + arguments = _arguments() + calls = [] + signature = build_signature(_vector_add, arguments) + + class Kernel: + def __init__(self, path, symbol, size): + calls.append(("load", Path(path).read_bytes(), symbol, size)) + + def launch(self, *args): + calls.append(("launch", args)) + + monkeypatch.setitem(sys.modules, "cuda.tile._backends.cpu._cpu", SimpleNamespace(Kernel=Kernel)) + + runtime = _runtime(_vector_add, arguments) + loaded = cpu.load(b"object", signature, cpu.Options(False)) + cpu.launch(loaded, None, (4,), runtime) + with ct.backend_options({"num_cpu_threads": 3}): + cpu.launch(loaded, None, (4,), runtime) + + assert calls[0] == ("load", b"object", signature.symbol, 9 * 8) + assert isinstance(loaded.kernel, Kernel) + launches = [call for call in calls if call[0] == "launch"] + assert len(launches) == 2 + assert launches[0][1][1:] == ((4, 1, 1), 0) + assert loaded.packer.unpack(launches[0][1][0]) == tuple( + cpu._flatten_arguments(loaded.parameters, runtime)) + assert [args[2] for _, args in launches] == [0, 3] + + +@pytest.mark.usefixtures("native_cpu") +@pytest.mark.parametrize("grid", [(-1,), (1 << 31,), (1 << 22,) * 3]) +def test_launch_rejects_overflowing_grids(grid): + out = np.zeros((1, 1, 1), dtype=np.int32) + compiled = _backend.compile_for_launch(_grid_ids, (out,)) + with pytest.raises(ValueError, match="grid"): + _backend.launch_compiled(None, grid, compiled, (out,)) + + +@pytest.mark.usefixtures("native_cpu") +@pytest.mark.parametrize("threads", [-1, 1 << 31, 1.5, "3"]) +def test_launch_rejects_invalid_thread_counts(threads): + out = np.zeros((1, 1, 1), dtype=np.int32) + compiled = _backend.compile_for_launch(_grid_ids, (out,)) + with ct.backend_options({"num_cpu_threads": threads}): + with pytest.raises((TypeError, ValueError)): + _backend.launch_compiled(None, (1,), compiled, (out,)) + + +def test_integer_scalars_do_not_silently_wrap(): + with pytest.raises(struct.error): + cpu._argument_packer(("i",)).pack(1 << 31) + with pytest.raises(struct.error): + cpu._argument_packer(("q",)).pack(1 << 63) + packer = cpu._argument_packer(("q",)) + assert packer.unpack(packer.pack(1 << 40)) == (1 << 40,) + + +def test_packed_argument_slots_have_fixed_offsets(): + packer = cpu._argument_packer(("?", "i", "f", "q")) + data = packer.pack(True, -3, 1.5, 1 << 40) + assert len(data) == 4 * 8 + assert struct.unpack_from("@?", data, 0) == (True,) + assert struct.unpack_from("@i", data, 8) == (-3,) + assert struct.unpack_from("@f", data, 16) == (1.5,) + assert struct.unpack_from("@q", data, 24) == (1 << 40,) + + +@pytest.fixture +def native_cpu(): + pytest.importorskip("cuda.tile._backends.cpu._cpu") + compiler = pytest.importorskip("cuda.tile._backends._compiler") + if compiler.backend != "cpu": + pytest.skip("requires the native CPU compiler") + _backend.set_backend("cpu") + yield + _backend.set_backend(None) + + +@pytest.mark.usefixtures("native_cpu") +def test_lowering_pipeline_expands_integer_powers(): + source = b""" +module { + func.func @cube(%base: f32) -> f32 { + %power = arith.constant 3 : i32 + %result = math.fpowi %base, %power : f32, i32 + return %result : f32 + } +} +""" + llvm_ir = cpu.native_compiler("cpu").lower_cpu_mlir( + source) + assert b"fmul float" in llvm_ir + assert b"llvm.powi" not in llvm_ir + assert b"define float @cube" in llvm_ir + + +@pytest.mark.usefixtures("native_cpu") +def test_lowering_pipeline_converts_loop_carried_amx_tiles(): + source = b""" +module { + func.func @amx_loop(%a: memref<16x32xbf16>, %b: memref<16x32xbf16>, + %c: memref<16x16xf32>, %n: index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %zero = x86.amx.tile_zero : !x86.amx.tile<16x16xf32> + %result = scf.for %i = %c0 to %n step %c1 + iter_args(%acc = %zero) -> (!x86.amx.tile<16x16xf32>) { + %lhs = x86.amx.tile_load %a[%c0, %c0] + : memref<16x32xbf16> into !x86.amx.tile<16x32xbf16> + %rhs = x86.amx.tile_load %b[%c0, %c0] + : memref<16x32xbf16> into !x86.amx.tile<16x32xbf16> + %next = x86.amx.tile_mulf %lhs, %rhs, %acc + : !x86.amx.tile<16x32xbf16>, !x86.amx.tile<16x32xbf16>, + !x86.amx.tile<16x16xf32> + scf.yield %next : !x86.amx.tile<16x16xf32> + } + x86.amx.tile_store %c[%c0, %c0], %result + : memref<16x16xf32>, !x86.amx.tile<16x16xf32> + return + } +} +""" + llvm_ir = cpu.native_compiler("cpu").lower_cpu_mlir( + source) + assert b"unrealized_conversion_cast" not in llvm_ir + assert b"phi x86_amx" in llvm_ir + assert b"llvm.x86.tdpbf16ps.internal" in llvm_ir + + +@pytest.mark.skipif( + sys.platform != "linux" or "amx_tile" not in cpu._host_cpu().get("flags", "").split(), + reason="requires Linux with AMX tile support", +) +def test_native_import_enables_amx_tile_data(): + pytest.importorskip("cuda.tile._backends.cpu._cpu") + assert cpu._amx_enabled() + subprocess.run([sys.executable, "-c", """ +import ctypes +import torch + +# Create an OpenMP worker pool before enabling process-wide AMX permission. +torch.ones(64, 64).sum() +import cuda.tile._backends.cpu._cpu + +libc = ctypes.CDLL(None, use_errno=True) +permissions = ctypes.c_ulong() +result = libc.syscall(ctypes.c_long(158), ctypes.c_long(0x1022), + ctypes.byref(permissions)) +assert result == 0, ctypes.get_errno() +assert permissions.value & (1 << 18) +"""], check=True) + + +@pytest.mark.usefixtures("native_cpu") +@pytest.mark.parametrize("dtype", [ + np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.uint32, + np.uint64, np.float16, np.float32, np.float64]) +def test_native_vector_add(dtype): + a = np.arange(140, dtype=dtype)[::2] + b = np.ones(70, dtype=dtype) + c = np.zeros(140, dtype=dtype)[::2] + with ct.backend_options({"num_cpu_threads": 3}): + ct.launch(None, (5,), _vector_add, (a, b, c, 16)) + np.testing.assert_array_equal(c, a + b) + + +@ct.kernel +def _gather_scatter_add(a, b, out, tile_x: ConstInt, tile_y: ConstInt): + x = ct.bid(0) * tile_x + ct.arange(tile_x, dtype=ct.int32) + y = ct.bid(1) * tile_y + ct.arange(tile_y, dtype=ct.int32) + indices = (x[:, None], y[None, :]) + value = ct.gather(a, indices) + ct.gather(b, indices) + ct.scatter(out, indices, value) + + +@pytest.mark.usefixtures("native_cpu") +@pytest.mark.parametrize("tile", [(1, 16), (2, 8)]) +def test_native_2d_gather_scatter_with_partial_tiles(tile): + a = np.arange(7 * 38, dtype=np.float32).reshape(7, 38)[:, ::2] + b = np.ones(a.shape, dtype=np.float32) + storage = np.full((7, 38), -1, dtype=np.float32) + out = storage[:, ::2] + grid = tuple((size + extent - 1) // extent for size, extent in zip(a.shape, tile)) + with ct.backend_options({"num_cpu_threads": 3}): + ct.launch(None, grid, _gather_scatter_add, (a, b, out, *tile)) + np.testing.assert_array_equal(out, a + b) + np.testing.assert_array_equal(storage[:, 1::2], -1) + + +@ct.kernel +def _grid_ids(out): + x, y, z = ct.bid(0), ct.bid(1), ct.bid(2) + value = x + 10 * y + 100 * z + value += ct.num_blocks(0) * 1000 + ct.num_blocks(1) * 10000 + value += ct.num_blocks(2) * 100000 + ct.store(out, index=(x, y, z), tile=ct.full((1, 1, 1), value, ct.int32)) + + +@pytest.mark.usefixtures("native_cpu") +@pytest.mark.parametrize("grid", [(1,), (4,), (3, 2), (3, 2, 2), (0, 2, 2)]) +@pytest.mark.parametrize("threads", [1, 3]) +def test_native_grid_ids_and_sizes(grid, threads): + shape = _backend._util.normalize_dims(grid) + out = np.full(shape, -1, dtype=np.int32) + with ct.backend_options({"num_cpu_threads": threads}): + ct.launch(None, grid, _grid_ids, (out,)) + x, y, z = np.indices(shape) + expected = x + y * 10 + z * 100 + shape[0] * 1000 + shape[1] * 10000 + shape[2] * 100000 + np.testing.assert_array_equal(out, expected) + + +@ct.kernel +def _scalar_store(out, flag: bool, scale: float, count: int): + value = scale * count if flag else -scale + ct.store(out, index=(0,), tile=ct.full((1,), value, ct.float32)) + + +@pytest.mark.usefixtures("native_cpu") +@pytest.mark.parametrize("flag", [True, False]) +def test_native_scalars(flag): + out = torch.empty(1, dtype=torch.float32) + ct.launch(None, (1,), _scalar_store, (out, flag, 1.5, -7)) + assert out.item() == (-10.5 if flag else -1.5) + + +@pytest.mark.usefixtures("native_cpu") +def test_native_repeated_torch_launches_use_current_arguments(monkeypatch): + validated = [] + flatten = cpu._flatten_arguments + monkeypatch.setattr(cpu, "_flatten_arguments", lambda *args: ( + validated.append(args), flatten(*args))[1]) + a = torch.empty(64, dtype=torch.float32) + b = torch.ones(64, dtype=torch.float32) + c = torch.empty(64, dtype=torch.float32) + for scale in range(3): + a.copy_(torch.arange(64, dtype=torch.float32) * scale) + with ct.backend_options({"num_cpu_threads": 3}): + ct.launch(None, (4,), _vector_add, (a, b, c, 16)) + torch.testing.assert_close(c, a + b) + assert len(validated) == 1 + + +@pytest.fixture +def launch_cache(native_cpu): + def create(kernel, args, limit=4): + compiled = _backend.compile_for_launch(kernel, args) + _backend.launch_compiled(None, (0,), compiled, args) + cache = cpu.make_launch_cache(constant_parameters(kernel), limit) + cache.insert(args, compiled.loaded) + return cache, compiled + + return create + + +def test_native_cache_reads_current_pointers_and_scalars(launch_cache): + args = (torch.empty(17)[1:], True, 1.5, -7) + cache, _ = launch_cache(_scalar_store, args) + for flag, scale, count in ((True, 2.5, 3), (False, -3.0, -4)): + out = torch.empty(17)[1:] + assert cache.launch((out, flag, scale, count), None, (1,), {}) + assert out[0].item() == (scale * count if flag else -scale) + + +def test_native_cache_rejects_changed_layouts_and_devices(launch_cache): + args = tuple(torch.empty(17)[1:] for _ in range(3)) + (16,) + cache, _ = launch_cache(_vector_add, args) + for replacement in (torch.empty(8), torch.empty(32)[::2], + torch.empty(16, dtype=torch.int32), + torch.empty(16, device="meta"), args[0].reshape(4, 4)): + assert cache.launch((replacement, *args[1:]), None, (1,), {}) is not True + assert len(cache) == 1 + + +@pytest.mark.parametrize("count", [-(1 << 31) - 1, 1 << 31, 1 << 100]) +def test_native_cache_falls_back_on_scalar_overflow(launch_cache, count): + args = (torch.empty(16), True, 1.5, 1) + cache, _ = launch_cache(_scalar_store, args) + assert not cache.launch((*args[:3], count), None, (1,), {}) + with pytest.raises((struct.error, OverflowError)): + ct.launch(None, (1,), _scalar_store, (*args[:3], count)) + + +@pytest.mark.parametrize("annotation,bits", [(int, 32), (ct.ScalarInt64, 64)]) +def test_native_cache_integer_packing_matches_abi(launch_cache, annotation, bits): + @ct.kernel + def store(out, value: annotation): + ct.store(out, index=(0,), tile=ct.full((1,), value, ct.int64)) + + low = -(1 << (bits - 1)) + high = (1 << (bits - 1)) - 1 + out = torch.empty(16, dtype=torch.int64) + cache, _ = launch_cache(store, (out, low)) + for value in (low, high): + assert cache.launch((out, value), None, (1,), {}) + assert out[0].item() == value + for value in (low - 1, high + 1): + assert not cache.launch((out, value), None, (1,), {}) + + +@pytest.mark.parametrize("grid", [(), (-1,), (1, 1, 1, 1), (1 << 31,)]) +def test_native_cache_validates_grid(launch_cache, grid): + args = (torch.empty(16), True, 1.5, 1) + cache, _ = launch_cache(_scalar_store, args) + with pytest.raises(ValueError): + cache.launch(args, None, grid, {}) + + +def test_native_cache_synchronizes_stream_and_checks_threads(launch_cache): + args = (torch.empty(16), True, 1.5, 1) + cache, _ = launch_cache(_scalar_store, args) + calls = [] + stream = SimpleNamespace(synchronize=lambda: calls.append("sync")) + assert cache.launch(args, stream, (1,), {"num_cpu_threads": 1}) + assert calls == ["sync"] + with pytest.raises(ValueError, match="num_cpu_threads"): + cache.launch(args, None, (1,), {"num_cpu_threads": -1}) + with pytest.raises(TypeError, match="synchronize"): + cache.launch(args, object(), (1,), {}) + + +def test_native_cache_supports_concurrent_launches(launch_cache): + args = (torch.empty(17)[1:], True, 1.5, 1) + cache, _ = launch_cache(_scalar_store, args) + + def run(value): + out = torch.empty(17)[1:] + assert cache.launch((out, True, 2.0, value), None, (1,), {}) + return out[0].item() + + with ThreadPoolExecutor(max_workers=4) as executor: + assert list(executor.map(run, range(16))) == [2.0 * value for value in range(16)] + + +@ct.kernel +def _constant_store(out, value: ct.Constant[float]): + ct.store(out, index=(0,), tile=ct.full((1,), value, ct.float32)) + + +@pytest.mark.usefixtures("native_cpu") +def test_native_constant_cache_preserves_signed_zero(): + out = np.zeros(1, dtype=np.float32) + for value in (0.0, -0.0, 0.0): + ct.launch(None, (1,), _constant_store, (out, value)) + assert np.signbit(out[0]) == np.signbit(value) + + +def test_native_cache_constants_and_lru(launch_cache): + out = torch.empty(16) + cache, _ = launch_cache(_constant_store, (out, 0.0), limit=2) + for value in (-0.0, 1.0): + args = (out, value) + compiled = _backend.compile_for_launch(_constant_store, args) + _backend.launch_compiled(None, (0,), compiled, args) + cache.insert(args, compiled.loaded) + assert len(cache) == 2 + assert isinstance(cache.launch((out, 0.0), None, (1,), {}), bytes) + assert cache.launch((out, -0.0), None, (1,), {}) + assert torch.signbit(out[0]) + assert cache.launch((out, 1.0), None, (1,), {}) + assert out[0].item() == 1.0 + + +def test_native_cache_accepts_numpy_and_parameter_exports(launch_cache): + left = np.arange(16, dtype=np.float32) + right = np.ones(16, dtype=np.float32) + out = np.empty_like(left) + args = (left, right, out, 16) + cache, _ = launch_cache(_vector_add, args) + assert cache.launch(args, None, (1,), {}) + np.testing.assert_array_equal(out, left + right) + readonly = left.copy() + readonly.flags.writeable = False + assert cache.launch((readonly, right, out, 16), None, (1,), {}) is None + + left = torch.nn.Parameter(torch.arange(16, dtype=torch.float32), + requires_grad=False) + right = torch.ones(16) + out = torch.empty_like(right) + args = (left, right, out, 16) + cache, _ = launch_cache(_vector_add, args) + assert cache.launch(args, None, (1,), {}) + torch.testing.assert_close(out, left + right) + + +@pytest.mark.usefixtures("native_cpu") +@pytest.mark.filterwarnings("ignore:Setting the dtype on a NumPy array:DeprecationWarning") +def test_native_cache_does_not_insert_mutated_launch_arguments(): + out = np.empty(16, dtype=np.float32) + stream = SimpleNamespace(synchronize=lambda: setattr(out, "dtype", np.dtype("int32"))) + ct.launch(stream, (0,), _constant_store, (out, 1.0)) + with pytest.raises(TypeCheckingError, match="float32 to int32"): + ct.launch(None, (1,), _constant_store, (out, 1.0)) + + +@pytest.mark.usefixtures("native_cpu") +def test_native_cache_discovers_parameter_after_torch_import(): + code = """ +import sys +from types import SimpleNamespace +from cuda.tile._backends.cpu._cpu import LaunchCache +assert "torch" not in sys.modules +cache = LaunchCache((False,), 4) +import torch +parameter = torch.nn.Parameter(torch.empty(16), requires_grad=False) +cache.insert((parameter,), SimpleNamespace(native_plan=((1, "i"),))) +assert len(cache) == 1 +""" + subprocess.run([sys.executable, "-c", code], check=True) + + +@pytest.mark.usefixtures("native_cpu") +def test_native_cache_falls_back_for_singleton_strides_and_subclasses(): + class Tensor(torch.Tensor): + pass + + cache = cpu.make_launch_cache((False, False, False, True), 2) + arrays = tuple(torch.empty(16) for _ in range(3)) + for first in (torch.empty_strided((1, 16), (64, 1)), arrays[0].as_subclass(Tensor)): + args = (first, *arrays[1:], 16) + cache.insert(args, SimpleNamespace(native_plan=())) + assert cache.launch(args, None, (1,), {}) is None + assert len(cache) == 0 + + +@ct.kernel +def _reduce_and_exp(a, out): + row = ct.bid(0) + values = ct.load(a, index=(row, 0), shape=(1, 16)) + result = ct.sum(ct.exp(values), axis=1) + ct.store(out, index=(row,), tile=result) + + +@pytest.mark.usefixtures("native_cpu") +def test_native_reduction_and_libm(): + a = torch.linspace(-1, 1, 64).reshape(4, 16) + out = torch.empty(4) + with ct.backend_options({"num_cpu_threads": 3}): + ct.launch(None, (4,), _reduce_and_exp, (a, out)) + torch.testing.assert_close(out, a.exp().sum(dim=1)) + + +@pytest.mark.usefixtures("native_cpu") +@pytest.mark.skipif( + sys.platform != "linux" or "amx_bf16" not in cpu._host_cpu().get("flags", "").split(), + reason="requires AMX BF16 support", +) +@pytest.mark.parametrize("batched", [False, True]) +def test_cpu_matmul_sample_emits_amx_bf16(monkeypatch, tmp_path, batched): + from cuda.tile._cext import default_tile_context + + if not cpu._amx_enabled(): + pytest.skip("AMX tile-data permission is unavailable") + monkeypatch.setattr(default_tile_context.config, "cache_dir", None) + dumps = [] + invoke = cpu.invoke + + def dumping_invoke(*args): + dump = tmp_path / f"{len(dumps)}.mlir" + dumps.append(dump) + return invoke(*args[:5], str(dump), *args[6:]) + + monkeypatch.setattr(cpu, "invoke", dumping_invoke) + filename = "BatchMatMul-cpu.py" if batched else "MatMul-cpu.py" + path = Path(__file__).resolve().parents[3] / "samples" / "cpu" / filename + spec = importlib.util.spec_from_file_location("cpu_sample_matmul_amx", path) + sample = importlib.util.module_from_spec(spec) + spec.loader.exec_module(sample) + shape = (2, 64, 64) if batched else (32, 32) + a = torch.randn(shape, dtype=torch.bfloat16) + b = torch.randn(shape, dtype=torch.bfloat16) + if batched: + out = sample.bmm(a, b, a.dtype, num_cpu_threads=1) + else: + out = sample.cutile_matmul(a, b, num_cpu_threads=1) + torch.testing.assert_close(out, a @ b, rtol=2e-2, atol=2e-2) + + assert dumps + assert cpu._transform_entry_point() == "__transform_nano_amx" + native = cpu.native_compiler("cpu") + source = dumps[0].read_bytes() + amx_ir = native.lower_cpu_mlir( + source, "__transform_nano_amx") + assert b"@llvm.x86.tdpbf16ps" in amx_ir + for entry_point in ("__transform_nano_x86", "__transform_nano"): + llvm_ir = native.lower_cpu_mlir( + source, entry_point) + assert b"@llvm.x86.tdpbf16ps" not in llvm_ir + + +@pytest.mark.usefixtures("native_cpu") +def test_amx_schedule_lowers_blocked_bfloat16_matmul(): + source = b""" +module { + func.func @blocked_matmul( + %a: memref>, + %b: memref>, + %out: memref<32x32xf32>, %m: index, %n: index, + %start: index, %end: index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %pad = arith.constant 0.0 : bf16 + %zero = arith.constant dense<0.0> : vector<32x32xf32> + %result = scf.for %k = %start to %end step %c1 + iter_args(%acc = %zero) -> (vector<32x32xf32>) { + %av = vector.transfer_read %a[%m, %k, %c0, %c0], %pad + {in_bounds = [true, true, true, true]} + : memref>, vector<1x1x32x32xbf16> + %lhs = vector.shape_cast %av : vector<1x1x32x32xbf16> to vector<32x32xbf16> + %bv = vector.transfer_read %b[%n, %k, %c0, %c0, %c0], %pad + {in_bounds = [true, true, true, true, true]} + : memref>, vector<1x1x16x32x2xbf16> + %bt = vector.transpose %bv, [0, 1, 2, 4, 3] + : vector<1x1x16x32x2xbf16> to vector<1x1x16x2x32xbf16> + %rhs = vector.shape_cast %bt : vector<1x1x16x2x32xbf16> to vector<32x32xbf16> + %next = vector.contract { + indexing_maps = [affine_map<(m, n, k) -> (m, k)>, + affine_map<(m, n, k) -> (k, n)>, + affine_map<(m, n, k) -> (m, n)>], + iterator_types = ["parallel", "parallel", "reduction"], + kind = #vector.kind} + %lhs, %rhs, %acc + : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32> + scf.yield %next : vector<32x32xf32> + } + vector.transfer_write %result, %out[%c0, %c0] {in_bounds = [true, true]} + : vector<32x32xf32>, memref<32x32xf32> + return + } +} +""" + llvm_ir = cpu.native_compiler("cpu").lower_cpu_mlir( + source, "__transform_nano_amx") + dot_products = [line for line in llvm_ir.splitlines() + if b"call " in line and b"@llvm.x86.tdpbf16ps" in line] + assert len(dot_products) == 4 + assert b"@llvm.fma" not in llvm_ir + + +@pytest.mark.usefixtures("native_cpu") +@pytest.mark.parametrize("scaling_block_size", [16, 32]) +def test_swizzled_scale_tile_shape_is_independent_of_scaling_block_size(scaling_block_size): + path = Path(__file__).resolve().parents[3] / "samples" / "cpu" / "BlockScaledMatMul-cpu.py" + spec = importlib.util.spec_from_file_location("cpu_sample_block_scaled_matmul", path) + sample = importlib.util.module_from_spec(spec) + spec.loader.exec_module(sample) + k = 4 * scaling_block_size + scales = sample.swizzle_32_4_4(torch.ones((128, 4), dtype=torch.bfloat16)) + args = (torch.ones((128, k), dtype=torch.bfloat16), scales, + torch.ones((k, 128), dtype=torch.bfloat16), scales, + torch.empty((128, 128)), 128, 128, k, scaling_block_size) + kernel = sample.block_scaled_matmul_kernel + bytecode, _ = _custom.compile_bytecode(kernel, cpu, build_signature(kernel, args)) + assert bytecode + + +@ct.kernel +def _wide_scalar(out: ct.IndexedWithInt64, value: ct.ScalarInt64): + ct.store(out, index=(0,), tile=ct.full((1,), value, ct.int64)) + + +@pytest.mark.usefixtures("native_cpu") +def test_native_int64_scalar_and_array_metadata(): + out = np.zeros(1, dtype=np.int64) + value = (1 << 40) + 37 + ct.launch(None, (1,), _wide_scalar, (out, value)) + assert out[0] == value + + +@ct.kernel +def _copy_static_shape(x: Annotated[ct.Array, ct.ArrayAnnotation(static_shape_dims=(0,))], + y, size: ConstInt): + ct.store(y, index=(0,), tile=ct.load(x, index=(0,), shape=(size,))) + + +@ct.kernel +def _copy_matrix(x, y, rows: ConstInt, cols: ConstInt): + ct.store(y, index=(0, 0), tile=ct.load(x, index=(0, 0), shape=(rows, cols))) + + +@pytest.mark.usefixtures("native_cpu") +def test_native_static_and_size_one_shapes(): + x = np.arange(16, dtype=np.float32) + y = np.zeros_like(x) + ct.launch(None, (1,), _copy_static_shape, (x, y, 16)) + np.testing.assert_array_equal(y, x) + + x, y = x.reshape(1, 16), np.zeros((1, 16), dtype=np.float32) + ct.launch(None, (1,), _copy_matrix, (x, y, 1, 16)) + np.testing.assert_array_equal(y, x) + + +@ct.kernel +def _empty_kernel(): + pass + + +@pytest.mark.usefixtures("native_cpu") +def test_native_kernel_without_arguments(): + with ct.backend_options({"num_cpu_threads": 3}): + ct.launch(None, (3, 2), _empty_kernel, ()) + + +@pytest.mark.usefixtures("native_cpu") +def test_existing_flat_cpu_abi_still_executes(monkeypatch, tmp_path): + args = (np.arange(64, dtype=np.float32), np.ones(64, dtype=np.float32), + np.zeros(64, dtype=np.float32), 16) + bytecode, signature = _custom.compile_bytecode( + _vector_add, cpu, build_signature(_vector_add, args)) + binary = cpu.native_compiler("cpu").compile_cpu( + bytecode, False, + str(tmp_path), "", False, False) + path = tmp_path / "flat.so" + path.write_bytes(binary) + library = ctypes.CDLL(str(path)) + function = getattr(library, signature.symbol) + function.argtypes = [*([ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32] * 3), + *([ctypes.c_int32] * 6)] + function.restype = None + values = cpu._flatten_arguments(runtime_parameters("CPU", signature), + _runtime(_vector_add, args)) + for block in range(4): + function(*values, block, 0, 0, 4, 1, 1) + np.testing.assert_array_equal(args[2], args[0] + args[1]) + + +@pytest.fixture(scope="module") +def openmp_probe(tmp_path_factory): + native = pytest.importorskip("cuda.tile._backends.cpu._cpu") + path = tmp_path_factory.mktemp("cpu-probe") / "probe.so" + source = b""" +#include +#include + +void probe(const void *arguments, int32_t x, int32_t y, int32_t z, + int32_t nx, int32_t ny, int32_t nz) { + struct { int32_t *output; int (*thread_id)(void); int (*team_size)(void); } args; + memcpy(&args, arguments, sizeof(args)); + int64_t index = x + (int64_t)nx * (y + (int64_t)ny * z); + args.output[index * 2] = args.thread_id(); + args.output[index * 2 + 1] = args.team_size(); +} +""" + compiler = shutil.which(os.environ.get("CC", "cc")) + if compiler is None: + pytest.skip("building the OpenMP probe requires a C compiler") + subprocess.run([compiler, "-shared", "-fPIC", "-x", "c", "-", "-o", str(path)], + input=source, check=True, timeout=60) + runtime = ctypes.CDLL(native.__file__) + return native, path, runtime + + +@pytest.mark.parametrize("threads", [0, 1, 3]) +def test_native_openmp_thread_count_and_static_order(openmp_probe, threads): + native, path, runtime = openmp_probe + old_dynamic = runtime.omp_get_dynamic() + runtime.omp_set_dynamic(0) + try: + expected_threads = threads or runtime.omp_get_max_threads() + grid = (expected_threads * 2, 2, 2) + out = np.full((np.prod(grid), 2), -1, dtype=np.int32) + packer = cpu._argument_packer(("P", "P", "P")) + args = packer.pack(out.ctypes.data, + ctypes.cast(runtime.omp_get_thread_num, ctypes.c_void_p).value, + ctypes.cast(runtime.omp_get_num_threads, ctypes.c_void_p).value) + kernel = native.Kernel(str(path), "probe", packer.size) + kernel.launch(args, grid, threads) + np.testing.assert_array_equal(out[:, 0], np.repeat(np.arange(expected_threads), 8)) + np.testing.assert_array_equal(out[:, 1], expected_threads) + finally: + runtime.omp_set_dynamic(old_dynamic) + + +def test_native_loader_reports_errors(openmp_probe): + native, path, _ = openmp_probe + with pytest.raises(RuntimeError, match="unable to load kernel"): + native.Kernel(str(path) + ".missing", "probe", 24) + with pytest.raises(RuntimeError, match="unable to find kernel"): + native.Kernel(str(path), "missing_symbol", 24) + with pytest.raises(ValueError, match="multiple of 8"): + native.Kernel(str(path), "probe", 3) + kernel = native.Kernel(str(path), "probe", 24) + with pytest.raises(ValueError, match="buffer size"): + kernel.launch(b"", (1, 1, 1), 1) diff --git a/backends/cpu/test/test_cpu_benchmark.py b/backends/cpu/test/test_cpu_benchmark.py new file mode 100644 index 00000000..9c4a1406 --- /dev/null +++ b/backends/cpu/test/test_cpu_benchmark.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import gc +import importlib.util +from pathlib import Path + +import pytest + + +def _load_benchmark(): + path = (Path(__file__).resolve().parents[3] + / "samples" / "cpu" / "utils" / "benchmark.py") + spec = importlib.util.spec_from_file_location("cpu_sample_benchmark", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +benchmark = _load_benchmark() + + +def _clock(values): + values = iter(values) + return lambda: next(values) + + +def test_calibration_targets_two_hundred_milliseconds(monkeypatch): + calls = 0 + + def function(): + nonlocal calls + calls += 1 + + monkeypatch.setattr( + benchmark.time, "perf_counter_ns", _clock((0, 5_000_000))) + + warmup, iterations, rounds = benchmark._estimate_bench_iter(function, ()) + + assert (warmup, iterations, rounds) == (1, 200, benchmark._ROUNDS) + assert calls == 6 + + +def test_calibration_caps_iteration_count(monkeypatch): + monkeypatch.setattr( + benchmark.time, "perf_counter_ns", _clock((0, 1))) + + _, iterations, _ = benchmark._estimate_bench_iter(lambda: None, ()) + + assert iterations == benchmark._MAX_ITERS + + +def test_calibration_rejects_stalled_timer(monkeypatch): + monkeypatch.setattr( + benchmark.time, "perf_counter_ns", _clock((7, 7))) + + with pytest.raises(RuntimeError, match="timer did not advance"): + benchmark._estimate_bench_iter(lambda: None, ()) + + +def test_report_benchmark_returns_arithmetic_mean(monkeypatch): + calls = 0 + + def function(): + nonlocal calls + calls += 1 + + rounds = benchmark._ROUNDS + clock_values = [0, 5_000_000] + for _ in range(rounds): + clock_values.extend((0, 300_000_000)) + monkeypatch.setattr( + benchmark.time, "perf_counter_ns", _clock(clock_values)) + + result = benchmark.report_benchmark(function, ()) + + assert result == {"mean_time_ms": 1.5} + assert calls == 7 + 200 * rounds + + +def test_measurement_restores_enabled_gc(): + gc.enable() + + def fail(): + raise ValueError("expected") + + with pytest.raises(ValueError, match="expected"): + benchmark._measure_block_ns(fail, (), 1) + + assert gc.isenabled() + + +def test_measurement_preserves_disabled_gc(): + gc.disable() + try: + benchmark._measure_block_ns(lambda: None, (), 1) + assert not gc.isenabled() + finally: + gc.enable() diff --git a/backends/cpu/test/test_cpu_samples.py b/backends/cpu/test/test_cpu_samples.py new file mode 100644 index 00000000..0b9cbd03 --- /dev/null +++ b/backends/cpu/test/test_cpu_samples.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 +"""Run every CPU sample with its correctness check, like samples/test_samples.py for CUDA.""" + +import os +from pathlib import Path +import subprocess +import sys + +import pytest + +_SAMPLES = sorted((Path(__file__).resolve().parents[3] / "samples" / "cpu").glob("*-cpu.py")) + + +@pytest.fixture(scope="module", autouse=True) +def _native_cpu(): + pytest.importorskip("cuda.tile._backends.cpu._cpu") + compiler = pytest.importorskip("cuda.tile._backends._compiler") + if compiler.backend != "cpu": + pytest.skip("requires the native CPU compiler") + + +@pytest.mark.parametrize("sample", _SAMPLES, ids=lambda path: path.name) +def test_cpu_sample_passes_correctness_check(sample): + # As in conftest.py, compile without pytest.ini's per-compile timeout, which is + # meant for tileiras: native CPU compiles of large tiles take minutes. + environment = dict(os.environ) + environment.pop("CUDA_TILE_COMPILER_TIMEOUT_SEC", None) + result = subprocess.run([sys.executable, "-u", str(sample), "--correctness-check"], + capture_output=True, text=True, timeout=1800, env=environment) + output = result.stdout + result.stderr + assert result.returncode == 0, output + assert "Correctness check passed" in result.stdout, output + assert "Correctness check disabled" not in result.stdout, output diff --git a/backends/cpu/toolchain.cmake b/backends/cpu/toolchain.cmake new file mode 100644 index 00000000..d85d0d05 --- /dev/null +++ b/backends/cpu/toolchain.cmake @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +set(_cutile_llvm_targets host) +set(_cutile_llvm_projects "mlir|lld") +set(_cutile_llvm_experimental_targets "") +set(_cutile_disable_ocloc ON) +set(_cutile_backend_cmake_args "") \ No newline at end of file diff --git a/backends/launch.h b/backends/launch.h new file mode 100644 index 00000000..3a506643 --- /dev/null +++ b/backends/launch.h @@ -0,0 +1,477 @@ +// SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cutile::launch { +namespace nb = nanobind; + +struct ManagedTensor { + nb::dlpack::dltensor tensor; + void *context; + void (*deleter)(ManagedTensor *); +}; + +struct Argument { + nb::handle scalar; + nb::object capsule; + const nb::dlpack::dltensor *tensor = nullptr; +}; + +struct Slot { + uint64_t value = 0; + size_t size = 8; +}; + +inline int64_t index(nb::handle value) { + nb::object integer = nb::steal(PyNumber_Index(value.ptr())); + if (!integer.is_valid()) + throw nb::python_error(); + int64_t result = PyLong_AsLongLong(integer.ptr()); + if (PyErr_Occurred()) + throw nb::python_error(); + return result; +} + +inline std::array dimensions(nb::handle grid) { + std::array result{1, 1, 1}; + if (PyLong_Check(grid.ptr())) { + result[0] = index(grid); + return result; + } + nb::object sequence = nb::steal(PySequence_Tuple(grid.ptr())); + if (!sequence.is_valid()) + throw nb::python_error(); + Py_ssize_t size = PySequence_Fast_GET_SIZE(sequence.ptr()); + if (size < 1 || size > 3) + throw nb::value_error("launch dimensions must have length 1 to 3"); + for (Py_ssize_t dim = 0; dim < size; ++dim) + result[dim] = + index(nb::handle(PySequence_Fast_GET_ITEM(sequence.ptr(), dim))); + return result; +} + +inline void synchronize(nb::handle stream) { + if (stream.is_none()) + return; + if (PyLong_Check(stream.ptr())) { + int zero = PyObject_RichCompareBool(stream.ptr(), nb::int_(0).ptr(), Py_EQ); + if (zero < 0) + throw nb::python_error(); + if (zero) + return; + } + nb::object function = + nb::steal(PyObject_GetAttrString(stream.ptr(), "synchronize")); + if (!function.is_valid()) { + if (!PyErr_ExceptionMatches(PyExc_AttributeError)) + throw nb::python_error(); + PyErr_Clear(); + } + if (!function.is_valid() || function.is_none()) + throw nb::type_error("stream must be None, 0, or provide synchronize()"); + function(); +} + +template void append(std::string &key, T value) { + key.append(reinterpret_cast(&value), sizeof(value)); +} + +template Slot slot(T value) { + Slot result; + result.size = sizeof(T); + std::memcpy(&result.value, &value, sizeof(value)); + return result; +} + +template bool signedSlot(int64_t number, Slot &result) { + if (number < std::numeric_limits::min() || + number > std::numeric_limits::max()) + return false; + result = slot(static_cast(number)); + return true; +} + +template bool unsignedSlot(uint64_t number, Slot &result) { + if (number > std::numeric_limits::max()) + return false; + result = slot(static_cast(number)); + return true; +} + +inline bool scalarSlot(nb::handle value, char format, Slot &result) { + if (format == '?') { + result = slot(bool(PyObject_IsTrue(value.ptr()))); + return true; + } + if (format == 'f' || format == 'd') { + double number = PyFloat_AsDouble(value.ptr()); + if (PyErr_Occurred()) { + PyErr_Clear(); + return false; + } + if (format == 'f') { + if (std::isfinite(number) && + std::abs(number) > std::numeric_limits::max()) + return false; + result = slot(static_cast(number)); + } else { + result = slot(number); + } + return true; + } + bool isUnsigned = + format == 'B' || format == 'H' || format == 'I' || format == 'Q'; + if (isUnsigned) { + uint64_t number = PyLong_AsUnsignedLongLong(value.ptr()); + if (PyErr_Occurred()) { + PyErr_Clear(); + return false; + } + switch (format) { + case 'B': return unsignedSlot(number, result); + case 'H': return unsignedSlot(number, result); + case 'I': return unsignedSlot(number, result); + case 'Q': result = slot(number); return true; + default: break; + } + } + int64_t number = PyLong_AsLongLong(value.ptr()); + if (PyErr_Occurred()) { + PyErr_Clear(); + return false; + } + switch (format) { + case 'b': return signedSlot(number, result); + case 'h': return signedSlot(number, result); + case 'i': return signedSlot(number, result); + case 'q': result = slot(number); return true; + default: + throw nb::value_error("unsupported native scalar format"); + } +} + +struct Parameter { + int rank; + char format; +}; + +inline std::vector plan(nb::tuple descriptors) { + std::vector result; + for (nb::handle descriptor : descriptors) { + if (PyTuple_Check(descriptor.ptr())) { + nb::tuple array = nb::borrow(descriptor); + result.push_back({nb::cast(array[0]), nb::cast(array[1])}); + } else { + result.push_back({-1, nb::cast(descriptor)}); + } + } + return result; +} + +template +bool pack(const std::vector &arguments, + const std::vector ¶meters, + typename Kernel::Arguments &slots) { + if (arguments.size() != parameters.size()) + return false; + auto add = [&](Slot value) { + if constexpr (Kernel::memref) + slots.push_back(value); + else + slots.push_back(value.value); + }; + for (size_t arg = 0; arg < arguments.size(); ++arg) { + const auto &argument = arguments[arg]; + const auto ¶meter = parameters[arg]; + if (parameter.rank < 0) { + Slot value; + if (argument.tensor || + !scalarSlot(argument.scalar, parameter.format, value)) + return false; + add(value); + continue; + } + if (!argument.tensor || argument.tensor->ndim != parameter.rank) + return false; + const auto &tensor = *argument.tensor; + uintptr_t pointer = + reinterpret_cast(tensor.data) + tensor.byte_offset; + add(slot(pointer)); + if (Kernel::memref) { + add(slot(pointer)); + add(slot(int64_t(0))); + } + for (const int64_t *dims : {tensor.shape, tensor.strides}) { + for (int dim = 0; dim < tensor.ndim; ++dim) { + int64_t value = dims[dim]; + if (value < 0) + return false; + if (Kernel::memref || parameter.format == 'q') { + add(slot(value)); + } else { + if (value > std::numeric_limits::max()) + return false; + add(slot(static_cast(value))); + } + } + } + } + return true; +} + +template class Cache { + struct Entry { + nb::object loaded; + std::vector parameters; + size_t slots = 0; + }; + struct Cached { + std::shared_ptr entry; + std::list::iterator position; + }; + +public: + Cache(nb::tuple constants, size_t limit) : limit(limit) { + for (nb::handle constant : constants) + this->constants.push_back(nb::cast(constant)); + discoverExporters(); + } + + nb::object launch(nb::tuple args, nb::handle stream, nb::handle grid, + nb::handle options) { + std::string key; + std::vector arguments; + if (!describe(args, key, arguments)) + return nb::none(); + std::shared_ptr entry; + { + std::lock_guard lock(mutex); + auto found = entries.find(key); + if (found == entries.end()) + return nb::bytes(key.data(), key.size()); + entry = found->second.entry; + order.splice(order.end(), order, found->second.position); + } + typename Kernel::Arguments slots; + slots.reserve(entry->slots); + if (!pack(arguments, entry->parameters, slots)) + return nb::none(); + auto dims = dimensions(grid); + synchronize(stream); + auto &kernel = nb::cast(entry->loaded.attr("kernel")); + kernel.launchNative(slots, dims, options, entry->loaded); + return nb::bool_(true); + } + + void insert(nb::tuple args, nb::object loaded, nb::handle expectedKey) { + std::string key; + std::vector arguments; + if (!limit || !describe(args, key, arguments)) + return; + if (!expectedKey.is_none()) { + nb::bytes expected = nb::cast(expectedKey); + if (key.size() != expected.size() || + std::memcmp(key.data(), expected.data(), key.size()) != 0) + return; + } + auto entry = + std::make_shared(Entry{loaded, plan(loaded.attr("native_plan"))}); + for (const auto ¶meter : entry->parameters) + entry->slots += parameter.rank < 0 + ? 1 + : 2 * parameter.rank + (Kernel::memref ? 3 : 1); + std::lock_guard lock(mutex); + auto found = entries.find(key); + if (found != entries.end()) { + found->second.entry.swap(entry); + order.splice(order.end(), order, found->second.position); + return; + } + order.push_back(key); + entries.emplace(std::move(key), + Cached{std::move(entry), std::prev(order.end())}); + if (entries.size() > limit) { + entry = std::move(entries.at(order.front()).entry); + entries.erase(order.front()); + order.pop_front(); + } + } + + size_t size() { + std::lock_guard lock(mutex); + return entries.size(); + } + +private: + void discoverExporters() { + PyObject *modules = PyImport_GetModuleDict(); + PyObject *torchModule = PyDict_GetItemString(modules, "torch"); + if (torchModule && !tensorType.is_valid()) { + nb::object torch = nb::borrow(torchModule); + if (nb::hasattr(torch, "Tensor") && nb::hasattr(torch, "_C")) { + tensorType = torch.attr("Tensor"); + toDlpack = torch.attr("_C").attr("_to_dlpack"); + } + } + PyObject *parameterModule = + PyDict_GetItemString(modules, "torch.nn.parameter"); + if (parameterModule && !parameterType.is_valid()) { + nb::object module = nb::borrow(parameterModule); + if (nb::hasattr(module, "Parameter")) + parameterType = module.attr("Parameter"); + } + PyObject *numpyModule = PyDict_GetItemString(modules, "numpy"); + if (numpyModule && !numpyArrayType.is_valid()) { + nb::object module = nb::borrow(numpyModule); + if (nb::hasattr(module, "ndarray")) + numpyArrayType = module.attr("ndarray"); + } + } + + bool describe(nb::tuple args, std::string &key, + std::vector &arguments) { + if (args.size() != constants.size()) + return false; + key.reserve(args.size() * 64); + arguments.reserve(args.size()); + for (size_t arg = 0; arg < args.size(); ++arg) { + nb::handle value = args[arg]; + uint8_t kind = PyBool_Check(value.ptr()) ? 1 + : PyLong_CheckExact(value.ptr()) ? 2 + : PyFloat_CheckExact(value.ptr()) ? 3 + : 4; + append(key, kind); + if (constants[arg]) { + if (kind == 1 || kind == 2) { + int overflow = 0; + int64_t number = PyLong_AsLongLongAndOverflow(value.ptr(), &overflow); + if (overflow || PyErr_Occurred()) { + PyErr_Clear(); + return false; + } + append(key, number); + } else if (kind == 3) { + append(key, PyFloat_AS_DOUBLE(value.ptr())); + } else { + return false; + } + continue; + } + if (kind != 4) { + arguments.push_back({value, {}, nullptr}); + continue; + } + PyObject *valueType = reinterpret_cast(Py_TYPE(value.ptr())); + bool torchTensor = tensorType.is_valid() && valueType == tensorType.ptr(); + bool torchParameter = parameterType.is_valid() && + valueType == parameterType.ptr(); + bool numpyArray = numpyArrayType.is_valid() && + valueType == numpyArrayType.ptr(); + if (!torchTensor && !torchParameter && !numpyArray) { + nb::object typeModule = nb::steal(PyObject_GetAttrString( + reinterpret_cast(Py_TYPE(value.ptr())), "__module__")); + if (!typeModule.is_valid()) + PyErr_Clear(); + const char *moduleName = + typeModule.is_valid() && PyUnicode_Check(typeModule.ptr()) + ? PyUnicode_AsUTF8(typeModule.ptr()) + : nullptr; + if (PyErr_Occurred()) { + PyErr_Clear(); + moduleName = nullptr; + } + bool mayBeTorch = + (!tensorType.is_valid() || !parameterType.is_valid()) && + moduleName && (std::strcmp(moduleName, "torch") == 0 || + std::strcmp(moduleName, "torch.nn.parameter") == 0); + bool mayBeNumpy = !numpyArrayType.is_valid() && + moduleName && std::strcmp(moduleName, "numpy") == 0; + if (mayBeTorch || mayBeNumpy) { + discoverExporters(); + torchTensor = tensorType.is_valid() && valueType == tensorType.ptr(); + torchParameter = + parameterType.is_valid() && valueType == parameterType.ptr(); + numpyArray = + numpyArrayType.is_valid() && valueType == numpyArrayType.ptr(); + } + } + if (!torchTensor && !torchParameter && !numpyArray) + return false; + if (!numpyArray && !toDlpack.is_valid()) + return false; + nb::object capsule = nb::steal( + numpyArray ? PyObject_CallMethod(value.ptr(), "__dlpack__", "") + : PyObject_CallOneArg(toDlpack.ptr(), value.ptr())); + if (!capsule.is_valid()) { + PyErr_Clear(); + return false; + } + auto *managed = static_cast( + PyCapsule_GetPointer(capsule.ptr(), "dltensor")); + if (!managed) { + PyErr_Clear(); + return false; + } + const auto &tensor = managed->tensor; + if (tensor.device.device_type != Kernel::device || + tensor.device.device_id != 0 || tensor.ndim < 0 || + (tensor.ndim && (!tensor.shape || !tensor.strides))) + return false; + append(key, tensor.dtype.code); + append(key, tensor.dtype.bits); + append(key, tensor.dtype.lanes); + append(key, tensor.ndim); + uintptr_t pointer = + reinterpret_cast(tensor.data) + tensor.byte_offset; + append(key, pointer & -pointer); + for (int dim = 0; dim < tensor.ndim; ++dim) { + // DLPack exporters may rewrite strides of singleton dimensions. + if (tensor.shape[dim] <= 1 || tensor.strides[dim] < 0) + return false; + append(key, tensor.shape[dim]); + append(key, tensor.strides[dim]); + } + arguments.push_back({{}, std::move(capsule), &tensor}); + } + return true; + } + + size_t limit; + std::vector constants; + nb::object tensorType, parameterType, numpyArrayType, toDlpack; + std::mutex mutex; + std::list order; + std::unordered_map entries; +}; + +template void bindCache(nb::module_ &module) { + using NativeCache = Cache; + nb::class_(module, "LaunchCache") + .def(nb::init()) + .def("launch", &NativeCache::launch, nb::arg("args"), + nb::arg("stream").none(), nb::arg("grid"), nb::arg("options")) + .def("insert", &NativeCache::insert, nb::arg("args"), nb::arg("loaded"), + nb::arg("expected_key").none() = nb::none()) + .def("__len__", &NativeCache::size); +} + +} // namespace cutile::launch \ No newline at end of file diff --git a/backends/mlir/LowerVectorScan.cpp b/backends/mlir/LowerVectorScan.cpp new file mode 100644 index 00000000..32dc2dcf --- /dev/null +++ b/backends/mlir/LowerVectorScan.cpp @@ -0,0 +1,39 @@ +// SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights +// reserved. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "LowerVectorScan.h" + +#include + +#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +namespace cutile { +namespace { + +class LowerVectorScanPass + : public mlir::PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerVectorScanPass) + + void runOnOperation() override { + mlir::RewritePatternSet patterns(&getContext()); + mlir::vector::populateVectorScanLoweringPatterns(patterns); + if (mlir::failed( + mlir::applyPatternsGreedily(getOperation(), std::move(patterns)))) + signalPassFailure(); + } +}; + +} // namespace + +std::unique_ptr createLowerVectorScanPass() { + return std::make_unique(); +} + +} // namespace cutile diff --git a/backends/mlir/LowerVectorScan.h b/backends/mlir/LowerVectorScan.h new file mode 100644 index 00000000..a252dd7e --- /dev/null +++ b/backends/mlir/LowerVectorScan.h @@ -0,0 +1,17 @@ +// SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights +// reserved. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include + +namespace mlir { +class Pass; +} + +namespace cutile { +// Lowers vector.scan to target-independent vector and arithmetic operations. +std::unique_ptr createLowerVectorScanPass(); +} // namespace cutile diff --git a/backends/mlir/compiler.cpp b/backends/mlir/compiler.cpp new file mode 100644 index 00000000..d9963ca6 --- /dev/null +++ b/backends/mlir/compiler.cpp @@ -0,0 +1,192 @@ +// SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights +// reserved. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "compiler.h" +#include "cuda_tile/Bytecode/Reader/BytecodeReader.h" +#include "cuda_tile/Dialect/CudaTile/IR/Attributes.h" +#include "cuda_tile/Dialect/CudaTile/IR/Dialect.h" +#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h" +#include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h" +#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h" +#include "mlir/Conversion/IndexToLLVM/IndexToLLVM.h" +#include "mlir/Conversion/MathToLLVM/MathToLLVM.h" +#include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h" +#include "mlir/Conversion/TileIRToMLIR/Passes.h" +#include "mlir/Conversion/UBToLLVM/UBToLLVM.h" +#include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h" +#include "mlir/Dialect/Affine/IR/AffineDialect.h" +#include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h" +#include "mlir/Dialect/Func/Extensions/InlinerExtension.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/Index/IR/IndexDialect.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Math/IR/Math.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.h" +#include "mlir/Dialect/UB/IR/UBOps.h" +#include "mlir/Dialect/Vector/IR/ValueBoundsOpInterfaceImpl.h" +#include "mlir/Dialect/Vector/IR/VectorOps.h" +#include "mlir/Dialect/X86/Transforms.h" +#include "mlir/Dialect/X86/X86Dialect.h" +#include "mlir/Parser/Parser.h" +#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h" +#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h" +#include "mlir/Transforms/Passes.h" +#include "llvm/Config/llvm-config.h" +#include "llvm/Support/MemoryBuffer.h" +#include +#include + +namespace cutile { + +static mlir::DialectRegistry createRegistry() { + mlir::DialectRegistry registry; + registry.insert(); + mlir::func::registerInlinerExtension(registry); + mlir::affine::registerValueBoundsOpInterfaceExternalModels(registry); + mlir::arith::registerValueBoundsOpInterfaceExternalModels(registry); + mlir::memref::registerValueBoundsOpInterfaceExternalModels(registry); + mlir::scf::registerValueBoundsOpInterfaceExternalModels(registry); + mlir::vector::registerValueBoundsOpInterfaceExternalModels(registry); + mlir::arith::registerConvertArithToLLVMInterface(registry); + mlir::cf::registerConvertControlFlowToLLVMInterface(registry); + mlir::registerConvertFuncToLLVMInterface(registry); + mlir::index::registerConvertIndexToLLVMInterface(registry); + mlir::registerConvertMathToLLVMInterface(registry); + mlir::registerConvertMemRefToLLVMInterface(registry); + mlir::ub::registerConvertUBToLLVMInterface(registry); + mlir::vector::registerConvertVectorToLLVMInterface(registry); + mlir::registerConvertX86ToLLVMInterface(registry); + mlir::registerBuiltinDialectTranslation(registry); + mlir::registerLLVMDialectTranslation(registry); + registerBackendTranslations(registry); + return registry; +} + +Compilation::Compilation(bool printIR) + : context(createRegistry(), mlir::MLIRContext::Threading::DISABLED), + printIR(printIR || + [] { + const char *value = std::getenv("MLIR_ENABLE_DUMP"); + return value && llvm::StringRef(value) == "1"; + }()), + handler(&context, [this](mlir::Diagnostic &diagnostic) { + llvm::raw_string_ostream stream(diagnostics); + diagnostic.print(stream); + stream << '\n'; + return mlir::success(); + }) {} + +mlir::OwningOpRef +Compilation::readTileIR(llvm::StringRef bytecode) { + context.getOrLoadDialect(); + auto buffer = llvm::MemoryBuffer::getMemBufferCopy(bytecode, ""); + if (!mlir::cuda_tile::isTileIRBytecode(buffer->getMemBufferRef())) + throw std::invalid_argument("cuTile compiler: invalid TileIR bytecode"); + auto decoded = + mlir::cuda_tile::readBytecode(buffer->getMemBufferRef(), context); + if (!decoded) + fail("reading TileIR bytecode"); + mlir::OwningOpRef module = + mlir::ModuleOp::create(mlir::UnknownLoc::get(&context)); + module->getBody()->push_back(decoded.release().getOperation()); + return module; +} + +mlir::OwningOpRef +Compilation::readMLIR(llvm::StringRef source) { + auto module = mlir::parseSourceString(source, &context); + if (!module) + fail("parsing MLIR"); + return module; +} + +void Compilation::lowerTileIR( + mlir::ModuleOp module, llvm::StringRef options, + llvm::function_ref addBackendPasses) { + module.walk([](mlir::Operation *op) { + auto rounding = + op->getAttrOfType("rounding_mode"); + if (!rounding) + return; + using Mode = mlir::cuda_tile::RoundingMode; + switch (rounding.getValue()) { + case Mode::NEAREST_EVEN: + case Mode::APPROX: + case Mode::FULL: + case Mode::NEAREST_INT_TO_ZERO: + return; + default: + throw std::invalid_argument( + std::string(backendName) + " backend: unsupported rounding mode " + + mlir::cuda_tile::stringifyRoundingMode(rounding.getValue()).str() + + " on " + op->getName().getStringRef().str()); + } + }); + mlir::PassManager pm(&context); + auto onError = [&](const llvm::Twine &message) { + diagnostics += message.str() + "\n"; + return mlir::failure(); + }; + if (failed(mlir::buildTileIRToMLIRPipeline(pm, options, onError))) + fail("constructing the TileIR pipeline"); + addBackendPasses(pm); + // Translating TileIR source locations to LLVM IR crashes; drop them. + pm.addPass(mlir::createStripDebugInfoPass()); + run(module, pm); + // Nearest-even is arith's default, including for narrowing casts. + module.walk([](mlir::Operation *op) { + auto rounding = + op->getAttrOfType("roundingmode"); + if (rounding && + rounding.getValue() == mlir::arith::RoundingMode::to_nearest_even) + op->removeAttr("roundingmode"); + }); +} + +void Compilation::run(mlir::ModuleOp module, mlir::PassManager &pm) { + if (printIR) + pm.enableIRPrinting([](mlir::Pass *, mlir::Operation *) { return true; }, + [](mlir::Pass *, mlir::Operation *) { return false; }); + if (failed(pm.run(module))) + fail("running the pass pipeline"); + if (!diagnostics.empty()) { + llvm::errs() << diagnostics; + diagnostics.clear(); + } +} + +void Compilation::fail(llvm::StringRef stage) { + throw std::runtime_error("cuTile compiler: failed " + stage.str() + "\n" + + diagnostics); +} + +std::string printModule(mlir::ModuleOp module) { + std::string text; + llvm::raw_string_ostream stream(text); + module.print(stream); + return text; +} + +} // namespace cutile + +NB_MODULE(_compiler, module) { + cutile::registerBackendPasses(); + module.attr("backend") = cutile::backendName; + module.attr("llvm_version") = LLVM_VERSION_STRING; + cutile::bindBackend(module); +} diff --git a/backends/mlir/compiler.h b/backends/mlir/compiler.h new file mode 100644 index 00000000..1bafa3fc --- /dev/null +++ b/backends/mlir/compiler.h @@ -0,0 +1,49 @@ +// SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights +// reserved. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Diagnostics.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/Pass/PassManager.h" +#include "llvm/ADT/STLFunctionalExtras.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/raw_ostream.h" +#include +#include +#include + +namespace cutile { + +// Backend hooks, implemented by exactly one backend's compiler source. +extern const char *const backendName; +void registerBackendTranslations(mlir::DialectRegistry ®istry); +void registerBackendPasses(); +void bindBackend(nanobind::module_ &module); + +class Compilation { +public: + explicit Compilation(bool printIR = false); + mlir::OwningOpRef readTileIR(llvm::StringRef bytecode); + mlir::OwningOpRef readMLIR(llvm::StringRef source); + // Runs the TileIR pipeline configured by `options`, then backend-specific + // passes. + void + lowerTileIR(mlir::ModuleOp module, llvm::StringRef options, + llvm::function_ref addBackendPasses); + void run(mlir::ModuleOp module, mlir::PassManager &pm); + [[noreturn]] void fail(llvm::StringRef stage); + mlir::MLIRContext context; + std::string diagnostics; + +private: + bool printIR; + mlir::ScopedDiagnosticHandler handler; +}; + +std::string printModule(mlir::ModuleOp module); + +} // namespace cutile diff --git a/backends/test/test_backend_dispatch.py b/backends/test/test_backend_dispatch.py new file mode 100644 index 00000000..f56c4705 --- /dev/null +++ b/backends/test/test_backend_dispatch.py @@ -0,0 +1,775 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +import struct +import threading +import time +from types import SimpleNamespace + +import numpy as np +import pytest + +import cuda.tile as ct +from cuda.tile import _backends as _backend +from cuda.tile._backends import ArrayArgument, _custom +from cuda.tile._backends._signature import build_signature, check_argument + + +@ct.kernel +def _kernel(x, scale: float, value: ct.Constant[int]): + pass + + +class _Backend: + name = "test" + bytecode_version = "13.3" + + def __init__(self, identity="test:v1"): + self.identity = identity + self.calls = [] + + def normalize_options(self, options): + self.calls.append("normalize") + return tuple(sorted(options.items())) + + def compiler_identity(self, options): + return None if self.identity is None else f"{self.identity}:{options}" + + def compile(self, bytecode, signature, options): + self.calls.append("compile") + time.sleep(0.01) + return b"binary" + + def load(self, binary, signature, options): + self.calls.append("load") + return binary, signature.symbol, options + + def launch(self, loaded, stream, grid, args): + self.calls.append(("launch", loaded, grid)) + self.launched_args = args + + def synchronize(self, stream): + self.calls.append("synchronize") + + +@pytest.fixture +def frontend(monkeypatch): + """Replace the TileIR frontend and the tile context by recording fakes.""" + from cuda.tile import _compile + + calls = [] + + def compile_tile(_function, signatures, *_args, **_kwargs): + calls.append("tileir") + return SimpleNamespace(kernel_signatures=list(signatures), + bytecode=b"tileir") + + monkeypatch.setattr(_compile, "compile_tile", compile_tile) + config = SimpleNamespace(cache_dir=None, cache_size_limit=1 << 20) + monkeypatch.setattr(_custom, "default_tile_context", + SimpleNamespace(config=config)) + return SimpleNamespace(calls=calls, config=config) + + +@pytest.fixture +def backend(): + backend = _Backend() + _backend.set_backend(backend) + yield backend + _backend.set_backend(None) + _kernel.__dict__.pop("_cutile_backend_cache", None) + + +def _args(scale=1.0, value=4): + return np.zeros(16, dtype=np.float32), scale, value + + +def test_launch_compiles_and_loads_once(frontend, backend): + x, _, _ = _args() + ct.launch(None, (1,), _kernel, (x, 1.0, 4)) + ct.launch(None, (2,), _kernel, (x, 2.0, 4)) + + launches = [call for call in backend.calls if call[0] == "launch"] + assert frontend.calls == ["tileir"] + assert backend.calls.count("compile") == 1 + assert backend.calls.count("load") == 1 + assert [grid for _, _, grid in launches] == [(1,), (2,)] + assert launches[0][1] == (b"binary", build_signature(_kernel, _args()).symbol, ()) + + +def test_launch_describes_each_array_once_and_drops_constants( + frontend, backend, monkeypatch +): + first = np.zeros(16, dtype=np.float32) + second = np.ones(16, dtype=np.float32) + described = [] + original_array_metadata = _custom.array_metadata + + def count_array_metadata(value): + described.append(value) + return original_array_metadata(value) + + monkeypatch.setattr(_custom, "array_metadata", count_array_metadata) + + ct.launch(None, (1,), _kernel, (first, 1.0, 4)) + first_args = backend.launched_args + ct.launch(None, (1,), _kernel, (second, 2.0, 4)) + + assert [value is first or value is second for value in described] == [True, False] * 2 + assert first_args == ( + ArrayArgument(first.ctypes.data, (16,), (1,), first.dtype, "cpu", 0), 1.0) + assert backend.launched_args[0].pointer == second.ctypes.data + assert backend.launched_args[1] == 2.0 + assert backend.calls.count("compile") == 1 + + +def test_launch_compiled_passes_runtime_arguments(frontend, backend): + x, _, _ = _args() + compiled = _backend.compile_for_launch(_kernel, (x, 1.0, 4)) + + _backend.launch_compiled(None, (3,), compiled, (x, 2.0, 4)) + + assert backend.calls[-1][2] == (3,) + assert backend.launched_args == (_custom.array_metadata(x), 2.0) + with pytest.raises(TypeError, match="expects 3 arguments"): + _backend.launch_compiled(None, (3,), compiled, (x, 2.0)) + + +def test_caches_are_bounded(frontend, backend, monkeypatch): + monkeypatch.setattr(_custom, "_CACHE_LIMIT", 2) + x, _, _ = _args() + + for value in range(4): + _backend.compile_for_launch(_kernel, (x, 1.0, value)) + + cache = _kernel._cutile_backend_cache + assert len(cache.launches) == len(cache.binaries) == 2 + _backend.compile_for_launch(_kernel, (x, 1.0, 3)) + assert backend.calls.count("compile") == 4 + _backend.compile_for_launch(_kernel, (x, 1.0, 0)) + assert backend.calls.count("compile") == 5 + + +class _HostTensor: + dtype = np.dtype("float32") + + def __init__(self, pointer=4096, shape=(16,), device="cpu"): + self.pointer, self.shape, self.device = pointer, shape, device + + def data_ptr(self): + return self.pointer + + def stride(self): + return (1,) * len(self.shape) + + +@pytest.fixture +def flat_backend(backend): + def launch(loaded, stream, grid, args): + signature = build_signature(_kernel, _args()) + check_argument("test", "cpu", signature.parameters[0], args[0]) + backend.calls.append("launch") + + def launch_flat(loaded, stream, grid, values): + backend.calls.append(("launch_flat", tuple(values))) + + backend.launch = launch + backend.launch_flat = launch_flat + return backend + + +def _launch_kinds(backend): + return [call if call == "launch" else call[0] for call in backend.calls + if call == "launch" or call[0] == "launch_flat"] + + +def test_launch_flat_reuses_only_accepted_layouts(frontend, flat_backend): + # 12288 has the same 4096-byte alignment as 4096; 4160 only has 64 bytes. + for args in [(_HostTensor(), 1.0, 4), (_HostTensor(12288), 2.0, 4), + (_HostTensor(shape=(8,)), 1.0, 4), (_HostTensor(4160), 1.0, 4), + (_HostTensor(), 1.0, 5), (_HostTensor(), np.float64(1.0), 4)]: + ct.launch(None, (1,), _kernel, args) + + assert _launch_kinds(flat_backend) == [ + "launch", "launch_flat", "launch", "launch", "launch", "launch"] + assert ("launch_flat", (12288, 16, 1, 2.0)) in flat_backend.calls + + +def test_launch_flat_never_caches_rejected_layouts(frontend, flat_backend): + for _ in range(2): + with pytest.raises(ValueError, match="expected 'cpu'"): + ct.launch(None, (1,), _kernel, (_HostTensor(device="cuda"), 1.0, 4)) + assert _launch_kinds(flat_backend) == [] + + +def test_launch_flat_normalizes_integer_metadata(frontend, flat_backend): + class Tensor(_HostTensor): + def stride(self): + return (1.0,) + + for tensor in (_HostTensor(), Tensor(shape=(16.0,)), Tensor(shape=(16.0,))): + ct.launch(None, (1,), _kernel, (tensor, 1.0, 4)) + + assert _launch_kinds(flat_backend) == ["launch", "launch_flat", "launch_flat"] + for call in flat_backend.calls: + if call[0] == "launch_flat": + assert tuple(map(type, call[1])) == (int, int, int, float) + + +def test_launch_flat_distinguishes_missing_and_none_device(frontend, flat_backend): + tensor = _HostTensor() + del tensor.device + ct.launch(None, (1,), _kernel, (tensor, 1.0, 4)) + ct.launch(None, (1,), _kernel, (tensor, 1.0, 4)) + + for _ in range(2): + with pytest.raises(ValueError, match="expected 'cpu'"): + ct.launch(None, (1,), _kernel, (_HostTensor(device=None), 1.0, 4)) + assert _launch_kinds(flat_backend) == ["launch", "launch_flat"] + + +def test_launch_flat_falls_back_for_other_arrays(frontend, flat_backend): + unhashable = SimpleNamespace(type="cpu", index=0) + for args in [_args(), _args(), (_HostTensor(device=unhashable), 1.0, 4), + (_HostTensor(device=unhashable), 1.0, 4)]: + ct.launch(None, (1,), _kernel, args) + assert _launch_kinds(flat_backend) == ["launch"] * 4 + + +def test_native_cache_unsupported_layout_does_not_insert(frontend, backend): + class NativeCache: + def launch(self, *args): + return None + + def insert(self, *args): + pytest.fail("unsupported layout must not be inserted") + + backend.make_launch_cache = lambda constants, limit: NativeCache() + x, _, _ = _args() + ct.launch(None, (1,), _kernel, (x, 1.0, 4)) + assert [call[0] for call in backend.calls if isinstance(call, tuple)] == ["launch"] + + +def test_native_caches_reuse_alternate_options(frontend, backend): + class NativeCache: + def __init__(self): + self.lookups = 0 + + def launch(self, *args): + self.lookups += 1 + return False + + def insert(self, *args): + pass + + caches = [] + + def make_cache(constants, limit): + cache = NativeCache() + caches.append(cache) + return cache + + backend.make_launch_cache = make_cache + x, _, _ = _args() + for option in (1, 2, 1): + with ct.backend_options({"variant": option}): + ct.launch(None, (1,), _kernel, (x, 1.0, 4)) + + assert len(caches) == 2 + assert [cache.lookups for cache in caches] == [2, 1] + + +def test_native_cache_memo_is_a_consistent_snapshot(backend): + parked = threading.Event() + resume = threading.Event() + launched = [] + + class Options: + def __init__(self, variant): + self.variant = variant + + def __hash__(self): + return hash(self.variant) + + def __eq__(self, other): + if threading.current_thread() is not threading.main_thread(): + parked.set() + assert resume.wait(5) + return self.variant == other.variant + + class NativeCache: + def __init__(self, variant): + self.variant = variant + + def launch(self, args, stream, grid, options): + launched.append((options["variant"], self.variant)) + return True + + backend.normalize_options = lambda options: Options(options["variant"]) + backend.make_launch_cache = lambda *args: NativeCache( + _custom.current_options()["variant"]) + cache = _custom._kernel_cache(_kernel) + first, second = NativeCache(1), NativeCache(2) + key = (backend, Options(1)) + cache.native_launch = (key, first) + cache.native_launches[key] = first + cache.native_launches[(backend, Options(2))] = second + + def launch_first(): + with ct.backend_options({"variant": 1}): + ct.launch(None, (0,), _kernel, _args()) + + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(launch_first) + try: + assert parked.wait(5) + with ct.backend_options({"variant": 2}): + ct.launch(None, (0,), _kernel, _args()) + finally: + resume.set() + future.result(timeout=5) + + assert launched == [(2, 2), (1, 1)] + + +def test_runtime_scalar_values_do_not_rebuild_signatures(frontend, backend): + built = [] + + def builder(kernel, args): + built.append(args[1]) + return build_signature(kernel, args) + + x, _, _ = _args() + first = _backend.compile_for_launch(_kernel, (x, 1.0, 4), signature_builder=builder) + second = _backend.compile_for_launch(_kernel, (x, 2.0, 4), signature_builder=builder) + + assert first is second + assert built == [1.0] + + +def test_signature_builder_tracks_alignment_above_4096(frontend, backend): + storage = np.zeros(32768, dtype=np.uint8) + offset = -storage.ctypes.data % 8192 + aligned = storage[offset:offset + 16] + less_aligned = storage[offset + 4096:offset + 4112] + + def builder(kernel, args): + return build_signature(kernel, args, base_addr_divisible_by=8192) + + first = _backend.compile_for_launch( + _kernel, (aligned, 1.0, 4), signature_builder=builder) + second = _backend.compile_for_launch( + _kernel, (less_aligned, 1.0, 4), signature_builder=builder) + assert first is not second + assert first.signature.parameters[0].base_addr_divisible_by == 8192 + assert second.signature.parameters[0].base_addr_divisible_by == 1 + assert _backend.compile_for_launch( + _kernel, (aligned, 2.0, 4), signature_builder=builder) is first + + +def test_signature_builder_tracks_non_power_of_two_alignment(frontend, backend): + class Tensor: + shape, dtype, device = (16,), np.dtype("uint8"), "cpu" + + def __init__(self, pointer): + self.pointer = pointer + + def data_ptr(self): + return self.pointer + + def stride(self): + return (1,) + + def builder(kernel, args): + return build_signature(kernel, args, base_addr_divisible_by=3) + + first = _backend.compile_for_launch( + _kernel, (Tensor(12288), 1.0, 4), signature_builder=builder) + second = _backend.compile_for_launch( + _kernel, (Tensor(20480), 1.0, 4), signature_builder=builder) + assert first.signature.parameters[0].base_addr_divisible_by == 3 + assert second.signature.parameters[0].base_addr_divisible_by == 1 + + +def test_constants_and_layouts_select_distinct_binaries(frontend, backend): + x = np.zeros(32, dtype=np.float32) + + by_constant = _backend.compile_for_launch(_kernel, (x, 1.0, 8)) + by_layout = _backend.compile_for_launch(_kernel, (x[::2], 1.0, 4)) + base = _backend.compile_for_launch(_kernel, (x, 1.0, 4)) + + assert len({id(by_constant), id(by_layout), id(base)}) == 3 + assert backend.calls.count("compile") == 3 + + +@pytest.mark.parametrize("values", [ + (0.0, -0.0), + (True, 1, 1.0), + tuple(struct.unpack("=d", struct.pack("=Q", bits))[0] + for bits in (0x7ff8000000000001, 0x7ff8000000000002)), +]) +def test_constant_cache_preserves_types_and_float_bits(frontend, backend, values): + x, _, _ = _args() + compiled = [_backend.compile_for_launch(_kernel, (x, 1.0, value)) + for value in values] + assert len({id(kernel) for kernel in compiled}) == len(values) + for value, kernel in zip(values, compiled): + assert _backend.compile_for_launch(_kernel, (x, 2.0, value)) is kernel + assert backend.calls.count("compile") == len(values) + + +def test_same_signature_shares_the_loaded_binary(frontend, backend): + x = np.zeros(64, dtype=np.float32) + + first = _backend.compile_for_launch(_kernel, (x[:16], 1.0, 4)) + second = _backend.compile_for_launch(_kernel, (x[16:32], 1.0, 4)) + + assert first is second + assert backend.calls.count("load") == 1 + + +def test_options_are_normalized_once_and_select_binaries(frontend, backend): + args = _args() + with ct.backend_options({"level": 1}): + first = _backend.compile_for_launch(_kernel, args) + again = _backend.compile_for_launch(_kernel, args) + with ct.backend_options({"level": 2}): + second = _backend.compile_for_launch(_kernel, args) + + assert first is again + assert first is not second + assert first.options == (("level", 1),) + assert backend.calls.count("normalize") == 2 + + +def test_disk_cache_reuses_backend_binary(frontend, backend, tmp_path): + frontend.config.cache_dir = str(tmp_path) + + _backend.compile_for_launch(_kernel, _args()) + _kernel.__dict__.pop("_cutile_backend_cache") + _backend.compile_for_launch(_kernel, _args()) + + assert frontend.calls == ["tileir", "tileir"] + assert backend.calls.count("compile") == 1 + assert backend.calls.count("load") == 2 + + +def test_unknown_identity_only_disables_the_disk_cache(frontend, backend, tmp_path): + backend.identity = None + frontend.config.cache_dir = str(tmp_path) + + first = _backend.compile_for_launch(_kernel, _args()) + again = _backend.compile_for_launch(_kernel, _args()) + _kernel.__dict__.pop("_cutile_backend_cache") + _backend.compile_for_launch(_kernel, _args()) + + assert first is again + assert backend.calls.count("compile") == 2 + + +def test_changed_identity_misses_the_disk_cache(frontend, backend, tmp_path): + frontend.config.cache_dir = str(tmp_path) + + _backend.compile_for_launch(_kernel, _args()) + _kernel.__dict__.pop("_cutile_backend_cache") + backend.identity = "test:v2" + _backend.compile_for_launch(_kernel, _args()) + + assert backend.calls.count("compile") == 2 + + +def test_explicit_symbols_do_not_share_binaries(frontend, backend): + def explicit(kernel, args): + return build_signature(kernel, args).with_symbol("explicit") + + default = _backend.compile_for_launch(_kernel, _args()) + renamed = _backend.compile_for_launch(_kernel, _args(), signature_builder=explicit) + + assert renamed.symbol == "explicit" + assert default.symbol != "explicit" + + +def test_concurrent_first_compile_runs_once(frontend, backend): + args = _args() + with ThreadPoolExecutor(max_workers=4) as executor: + results = list(executor.map( + lambda _: _backend.compile_for_launch(_kernel, args), range(4))) + + assert all(result is results[0] for result in results) + assert backend.calls.count("compile") == 1 + + +def test_benchmark_times_a_synchronized_launch(frontend, backend): + elapsed = _backend.benchmark(None, (1,), _kernel, _args()) + + assert elapsed >= 0 + assert backend.calls[-3:][0] == "synchronize" + assert backend.calls[-2][0] == "launch" + assert backend.calls[-1] == "synchronize" + + +def test_backend_protocol_lists_the_required_attributes(): + assert set(_custom._BACKEND_ATTRIBUTES) == { + "name", "bytecode_version", "normalize_options", "compiler_identity", + "compile", "load", "launch", "synchronize"} + + +def test_set_backend_rejects_incomplete_backends(): + with pytest.raises(TypeError, match="lacks .*load"): + _backend.set_backend(SimpleNamespace( + name="test", bytecode_version="13.3", + normalize_options=None, compiler_identity=None, compile=None, + launch=None, synchronize=None)) + assert _backend.get_backend() is None + + +@pytest.fixture +def cuda_launches(monkeypatch): + launches = [] + monkeypatch.setattr(_custom, "_cext_launch", + lambda *args, **options: launches.append((args, options))) + return launches + + +def test_launch_without_backend_uses_cuda(cuda_launches): + ct.launch(None, (1,), _kernel, _args(), programmatic_dependent_launch=True) + + assert len(cuda_launches) == 1 + assert cuda_launches[0][1] == {"programmatic_dependent_launch": True} + with pytest.raises(RuntimeError, match="no backend"): + _backend.compile_for_launch(_kernel, _args()) + + +class _CudaArray: + __cuda_array_interface__ = {} + + +class _CudaTensor: + shape, dtype = (16,), np.float32 + device = SimpleNamespace(type="cuda", index=0) + + def data_ptr(self): + return 256 + + def stride(self): + return (1,) + + +@dataclass +class _Pair: + left: object + right: object + + +def test_cuda_tensor_uses_selected_backend(frontend, backend, cuda_launches): + tensor = _CudaTensor() + ct.launch(None, (1,), _kernel, (tensor, 1.0, 4)) + + assert cuda_launches == [] + assert backend.launched_args == (_custom.array_metadata(tensor), 1.0) + assert backend.calls.count("compile") == 1 + + +def test_incompatible_device_raises_without_cuda_fallback( + frontend, backend, cuda_launches, monkeypatch): + def launch(loaded, stream, grid, args): + signature = build_signature(_kernel, _args()) + check_argument("test", "cpu", signature.parameters[0], args[0]) + + monkeypatch.setattr(backend, "launch", launch) + with pytest.raises(ValueError, match="device 'cuda', expected 'cpu'"): + ct.launch(None, (1,), _kernel, (_CudaTensor(), 1.0, 4)) + assert cuda_launches == [] + + +def test_explicit_backend_switching(frontend, backend, cuda_launches): + cpu_args = _args() + cuda_args = (_CudaTensor(), 1.0, 4) + ct.launch(None, (1,), _kernel, cpu_args) + + ct.set_backend(None) + assert _backend.get_backend() is None + ct.launch(None, (1,), _kernel, cuda_args, programmatic_dependent_launch=True) + + _backend.set_backend(backend) + ct.launch(None, (2,), _kernel, cpu_args) + + assert cuda_launches == [ + ((None, (1,), _kernel, cuda_args), {"programmatic_dependent_launch": True})] + assert backend.calls.count("compile") == 1 + assert [call[2] for call in backend.calls if call[0] == "launch"] == [(1,), (2,)] + + +@pytest.mark.parametrize("array", [_CudaArray(), _CudaTensor()]) +@pytest.mark.parametrize("wrap", [ + lambda array: array, + lambda array: (array,), + lambda array: [array], + lambda array: [(0, [array])], + lambda array: _Pair(0, (array,)), +]) +def test_cuda_arguments_do_not_enable_cuda_launch_options( + frontend, backend, cuda_launches, array, wrap): + args = (wrap(array), 1.0, 4) + with pytest.raises(TypeError, match="only supported on CUDA"): + ct.launch(None, (1,), _kernel, args, programmatic_dependent_launch=True) + assert cuda_launches == [] + assert "compile" not in backend.calls + + +def test_cuda_tuning_timeout_is_allowed_after_clearing_backend( + backend, monkeypatch): + from cuda.tile.tune import _tune + + calls = [] + monkeypatch.setattr(_tune, "_benchmark", lambda *args: 1.0) + monkeypatch.setattr(_tune, "benchmark_with_timeout", + lambda *args: calls.append(args)) + _backend.set_backend(None) + result = ct.tune.exhaustive_search( + [4], None, lambda _: (1,), _kernel, + lambda value: (_CudaTensor(), 1.0, value), + single_run_timeout_sec=1, quiet=True) + assert result.best.config == 4 + assert len(calls) == 1 + + +@pytest.mark.parametrize("array", [np.zeros(16, dtype=np.float32), _CudaTensor()]) +def test_backend_tuning_timeout_is_rejected(backend, array): + with pytest.raises(ValueError, match="not supported by custom backends"): + ct.tune.exhaustive_search( + [4], None, lambda _: (1,), _kernel, + lambda value: (array, 1.0, value), + single_run_timeout_sec=1, quiet=True) + + +def test_backends_reject_cuda_launch_options(frontend, backend): + with pytest.raises(TypeError, match="only supported on CUDA"): + ct.launch(None, (1,), _kernel, _args(), programmatic_dependent_launch=True) + + +@pytest.mark.parametrize("use_cuda", [False, True]) +def test_tuning_benchmark_uses_selected_backend(backend, monkeypatch, use_cuda): + from cuda.tile.tune import _tune + + calls = [] + monkeypatch.setattr(_backend, "benchmark", + lambda *args: (calls.append("backend") or 2.0)) + monkeypatch.setattr(_tune, "_benchmark", + lambda *args: (calls.append("cuda") or 1.0)) + if use_cuda: + _backend.set_backend(None) + + elapsed = _tune._benchmark_launch( + None, (1,), _kernel, (_CudaTensor(), 1.0, 4), False) + + assert calls == ["cuda" if use_cuda else "backend"] + assert elapsed == (1.0 if use_cuda else 2.0) + + +def test_autotune_prepare_follows_explicit_backend_switch( + frontend, backend, cuda_launches): + tuned = ct.autotune(configs=[{"value": 4}], grid=(1,))(_kernel) + args = (_CudaTensor(), 1.0) + + tuned.prepare(None, args) + tuned.launch_prepared(None) + assert backend.launched_args == (_custom.array_metadata(args[0]), 1.0) + assert cuda_launches == [] + + _backend.set_backend(None) + tuned.prepare(None, args) + tuned.launch_prepared(None) + assert len(cuda_launches) == 1 + assert cuda_launches[0][0][3] == (*args, 4) + + +def test_autotune_compiles_each_config_once(frontend, backend): + tuned = ct.autotune(configs=[{"value": 4}, {"value": 8}], grid=(1,))(_kernel) + x, _, _ = _args() + + tuned(None, (x, 1.0)) + tuned(None, (x, 2.0)) + tuned.prepare(None, (x, 3.0)) + tuned.launch_prepared(None) + + assert backend.calls.count("compile") == 2 + assert tuned.best_config.meta["value"] in (4, 8) + + +def test_autotune_retunes_when_backend_changes(backend, monkeypatch): + tuned = ct.autotune(configs=[{"value": 4}, {"value": 8}], grid=(1,))(_kernel) + calls = [] + + def tune(stream, args): + calls.append(_backend.get_backend()) + return len(calls) - 1, None + + monkeypatch.setattr(tuned, "_tune", tune) + args = _args()[:2] + assert tuned._resolve(None, args) == 0 + _backend.set_backend(_Backend("other")) + assert tuned._resolve(None, args) == 1 + _backend.set_backend(backend) + assert tuned._resolve(None, args) == 0 + assert len(calls) == 2 + + +@pytest.mark.parametrize("callable_options", [False, True]) +def test_autotune_retunes_when_launch_options_change(backend, monkeypatch, callable_options): + options = (lambda meta: {"level": meta["scale"]}) if callable_options else None + tuned = ct.autotune(configs=[{"value": 4}, {"value": 8}], + grid=(1,), options=options)(_kernel) + calls = [] + + def tune(stream, args): + calls.append(args) + return len(calls) - 1, None + + monkeypatch.setattr(tuned, "_tune", tune) + x = _args()[0] + for expected, value in ((0, 1), (1, 2), (0, 1)): + with ct.backend_options({"num_cpu_threads": value}): + assert tuned._resolve(None, (x, float(value))) == expected + assert len(calls) == 2 + + +def test_autotune_retunes_for_device_index(backend, monkeypatch): + tuned = ct.autotune(configs=[{"value": 4}, {"value": 8}], grid=(1,))(_kernel) + calls = [] + monkeypatch.setattr(tuned, "_tune", + lambda stream, args: (calls.append(args) or 0, None)) + tensor = _CudaTensor() + for index in (0, 1, 0): + tensor.device = SimpleNamespace(type="cuda", index=index) + tuned._resolve(None, (tensor, 1.0)) + assert len(calls) == 2 + + +def test_factory_autotune_retunes_when_environment_changes(backend, monkeypatch): + tuned = ct.autotune(configs=[{"level": 1}, {"level": 2}])( + lambda x, level=1: lambda: None) + calls = [] + monkeypatch.setattr(tuned, "_tune", + lambda args, kwargs: (calls.append(args) or 0, None)) + x = _args()[0] + for value in (1, 2, 1): + with ct.backend_options({"num_cpu_threads": value}): + tuned._resolve((x,), {}) + assert len(calls) == 2 + + +def test_autotune_uses_overridden_options_not_ambient_options(backend, monkeypatch): + tuned = ct.autotune(configs=[{"value": 4}, {"value": 8}], + grid=(1,), options={"level": 1})(_kernel) + calls = [] + monkeypatch.setattr(tuned, "_tune", + lambda stream, args: (calls.append(args) or 0, None)) + args = _args()[:2] + for value in (1, 2): + with ct.backend_options({"level": value}): + tuned._resolve(None, args) + assert len(calls) == 1 diff --git a/backends/test/test_compiler_build.py b/backends/test/test_compiler_build.py new file mode 100644 index 00000000..7a7a2bd9 --- /dev/null +++ b/backends/test/test_compiler_build.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +from pathlib import Path +import shutil +import subprocess +import tarfile + +import pytest + + +@pytest.fixture +def compiler_sources(tmp_path): + cmake = shutil.which("cmake") + if cmake is None: + pytest.skip("compiler fetch tests require CMake") + source = tmp_path / "project" + source.mkdir() + binary = tmp_path / "build" + fetch = Path(__file__).resolve().parents[1] / "cmake" / "FetchTileIRToMLIR.cmake" + (source / "CMakeLists.txt").write_text(f""" +cmake_minimum_required(VERSION 3.24) +project(CompilerFetchProbe NONE) +include("{fetch}") +fetch_tileir_to_mlir() +file(WRITE "${{CMAKE_BINARY_DIR}}/sources.txt" + "${{CUTILE_TILEIR_TO_MLIR_DIR}}\n${{CUTILE_CUDA_TILE_SOURCE_DIR}}\n") +""") + args = [] + archives = [] + for name, prefix, license_name in [ + ("tileir_to_mlir", "TILEIR_TO_MLIR", "LICENSE"), + ("cuda_tile", "CUDA_TILE", "LICENSE.txt"), + ]: + contents = tmp_path / name + contents.mkdir() + (contents / "CMakeLists.txt").write_text( + 'message(FATAL_ERROR "dependency must not be configured here")\n') + (contents / license_name).write_text(name) + (contents / "llvm-revision.txt").write_text("llvm-pin\n") + repository = tmp_path / f"{name}-repository" + (repository / "archive").mkdir(parents=True) + archive = repository / "archive" / "test-pin.tar.gz" + with tarfile.open(archive, "w:gz") as bundle: + bundle.add(contents, arcname=name) + archives.append(archive) + digest = hashlib.sha256(archive.read_bytes()).hexdigest() + args.extend([ + f"-D{prefix}_REPOSITORY={repository.as_uri()}", + f"-D{prefix}_REVISION=test-pin", + f"-D{prefix}_REVISION_SHA256={digest}", + ]) + + def run(extra_args=(), check=True): + return subprocess.run([cmake, "-S", str(source), "-B", str(binary), + *args, *extra_args], check=check, + capture_output=True, text=True, timeout=60) + + return run, binary, archives + + +def test_compiler_sources_are_fetched_and_reused(compiler_sources): + run, binary, archives = compiler_sources + run() + tileir, cuda_tile = map(Path, (binary / "sources.txt").read_text().splitlines()) + assert tileir == binary / "_deps" / "tileir_to_mlir-src" + assert cuda_tile == binary / "_deps" / "cuda_tile-src" + assert (tileir / "llvm-revision.txt").read_text() == "llvm-pin\n" + assert (tileir / "LICENSE").read_text() == "tileir_to_mlir" + assert (cuda_tile / "LICENSE.txt").read_text() == "cuda_tile" + for archive in archives: + archive.unlink() + run() + + +@pytest.mark.parametrize("prefix", ["TILEIR_TO_MLIR", "CUDA_TILE"]) +def test_compiler_source_hash_mismatch_is_rejected(compiler_sources, prefix): + run, _, _ = compiler_sources + result = run([f"-D{prefix}_REVISION_SHA256={'0' * 64}"], check=False) + assert result.returncode != 0 + assert "SHA256 hash" in result.stdout + result.stderr + + +@pytest.fixture +def llvm_build(tmp_path): + cmake, ninja = shutil.which("cmake"), shutil.which("ninja") + if cmake is None or ninja is None: + pytest.skip("compiler build tests require CMake and Ninja") + source = tmp_path / "source with spaces" + (source / "llvm").mkdir(parents=True) + cmakelists = source / "llvm" / "CMakeLists.txt" + cmakelists.write_text(""" +cmake_minimum_required(VERSION 3.20) +project(CompilerBuildProbe NONE) +if(TEST_VARIANT STREQUAL "fail-configure") + message(FATAL_ERROR "configuration failed") +endif() +file(APPEND "${CMAKE_BINARY_DIR}/configured.txt" "${TEST_VARIANT}:${TEST_LIST}\\n") +if(TEST_VARIANT STREQUAL "fail-build") + add_custom_target(probe COMMAND "${CMAKE_COMMAND}" -E false) +else() + add_custom_target(probe COMMAND "${CMAKE_COMMAND}" -E echo "Built ${TEST_VARIANT}") +endif() +""") + binary = tmp_path / "shared build" + template = (Path(__file__).resolve().parents[2] / "backends" / "cmake" + / "BuildLLVM.cmake.in").read_text() + + def run(variant, check=True): + script = tmp_path / f"build-{variant}.cmake" + contents = template + for key, value in { + "CUTILE_LLVM_SOURCE_DIR": str(source), + "CUTILE_LLVM_BINARY_DIR": str(binary), + "_cutile_llvm_cmake_args": f"-DTEST_VARIANT={variant};-DTEST_LIST=mlir|lld", + "_cutile_llvm_build_targets": "probe", + }.items(): + contents = contents.replace(f"@{key}@", value) + script.write_text(contents) + return subprocess.run([cmake, "-P", str(script)], check=check, + capture_output=True, text=True, timeout=60) + + return run, binary, cmakelists + + +def test_shared_build_only_configures_when_needed(llvm_build): + run, binary, cmakelists = llvm_build + run("cpu") + run("cpu") + log = binary / "configured.txt" + assert log.read_text().splitlines() == ["cpu:mlir;lld"] + run("other") + run("other") + run("cpu") + assert log.read_text().splitlines() == ["cpu:mlir;lld", "other:mlir;lld", "cpu:mlir;lld"] + # Manual cache edits must not defeat variant/configuration restoration. + cache = binary / "CMakeCache.txt" + cache.write_text(cache.read_text().replace( + "TEST_VARIANT:UNINITIALIZED=cpu", "TEST_VARIANT:UNINITIALIZED=manual")) + run("cpu") + assert log.read_text().splitlines()[-1] == "cpu:mlir;lld" + assert len(log.read_text().splitlines()) == 4 + # Ninja must still regenerate when a project's CMake sources change. + cmakelists.write_text(cmakelists.read_text() + '\nmessage(STATUS "source changed")\n') + run("cpu") + assert len(log.read_text().splitlines()) == 5 + run("cpu") + assert len(log.read_text().splitlines()) == 5 + + +@pytest.mark.parametrize("variant", ["fail-configure", "fail-build"]) +def test_failed_configuration_or_build_is_not_marked_successful(llvm_build, variant): + run, binary, _ = llvm_build + run("cpu") + result = run(variant, check=False) + assert result.returncode != 0 + assert not (binary / "cutile-configure.stamp").exists() + run("cpu") + assert (binary / "cutile-configure.stamp").exists() + + +def test_missing_cache_triggers_reconfiguration(llvm_build): + run, binary, _ = llvm_build + run("cpu") + (binary / "CMakeCache.txt").unlink() + run("cpu") + assert (binary / "configured.txt").read_text().splitlines() == [ + "cpu:mlir;lld", "cpu:mlir;lld"] + + +def test_cpu_toolchain_settings(tmp_path): + cmake = shutil.which("cmake") + if cmake is None: + pytest.skip("compiler build tests require CMake") + toolchain = Path(__file__).resolve().parents[1] / "cpu" / "toolchain.cmake" + script = tmp_path / "toolchain.cmake" + script.write_text(f'include("{toolchain}")\n' + """ +if(NOT _cutile_llvm_targets STREQUAL "host" OR + NOT _cutile_llvm_projects STREQUAL "mlir|lld" OR + NOT _cutile_llvm_experimental_targets STREQUAL "" OR + NOT _cutile_disable_ocloc OR + NOT _cutile_backend_cmake_args STREQUAL "") + message(FATAL_ERROR "CPU toolchain settings changed") +endif() +""") + subprocess.run([cmake, "-P", str(script)], check=True, + capture_output=True, text=True, timeout=60) diff --git a/backends/test/test_native_compiler.py b/backends/test/test_native_compiler.py new file mode 100644 index 00000000..3b35fb31 --- /dev/null +++ b/backends/test/test_native_compiler.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path +import subprocess +from types import SimpleNamespace + +import numpy as np +import pytest + +import cuda.tile as ct +from cuda.tile._backends import _native, cpu +from cuda.tile._cext import default_tile_context +from cuda.tile._exception import TileCompilerTimeoutError + + +def test_native_compiler_rejects_wrong_variant(monkeypatch): + monkeypatch.setattr( + _native.importlib, "import_module", lambda _: SimpleNamespace(backend="cpu")) + with pytest.raises(RuntimeError, match="targets cpu, not other"): + _native.compiler("other") + + +def test_default_compilation_stays_in_process(monkeypatch): + monkeypatch.setattr( + default_tile_context.config, "compiler_timeout_sec", None) + monkeypatch.setattr(_native, "compiler", lambda _: SimpleNamespace( + compile_cpu=lambda data, flag: data + bytes([flag]))) + monkeypatch.setattr( + _native.subprocess, "run", + lambda *args, **kwargs: pytest.fail("unexpected compiler subprocess")) + assert _native.invoke("cpu", "compile_cpu", b"input", True) == b"input\x01" + + +def test_explicit_timeout_uses_worker_and_reports_diagnostics(monkeypatch, capsys): + monkeypatch.setattr(default_tile_context.config, "compiler_timeout_sec", 7.5) + monkeypatch.setattr(_native, "compiler", lambda _: SimpleNamespace(__file__="/compiler")) + + def run(command, **kwargs): + assert Path(command[0]).is_absolute() + assert kwargs["timeout"] == 7.5 + assert kwargs["input"] == b"input" + return SimpleNamespace(returncode=0, stdout=b"binary", stderr=b"warning\n") + + monkeypatch.setattr(_native.subprocess, "run", run) + assert _native.invoke("cpu", "compile_cpu", b"input", True) == b"binary" + assert capsys.readouterr().err == "warning\n" + + +def test_native_timeout_is_reported(monkeypatch): + monkeypatch.setattr(default_tile_context.config, "compiler_timeout_sec", 0.1) + monkeypatch.setattr(_native, "compiler", lambda _: SimpleNamespace(__file__="/compiler")) + + def run(command, **kwargs): + raise subprocess.TimeoutExpired(command, kwargs["timeout"]) + + monkeypatch.setattr(_native.subprocess, "run", run) + with pytest.raises(TileCompilerTimeoutError, match="exceeded"): + _native.invoke("cpu", "compile_cpu", b"input") + + +def test_worker_failure_is_not_a_binary(monkeypatch): + monkeypatch.setattr(default_tile_context.config, "compiler_timeout_sec", 1) + monkeypatch.setattr(_native, "compiler", lambda _: SimpleNamespace(__file__="/compiler")) + monkeypatch.setattr(_native.subprocess, "run", lambda *args, **kwargs: + SimpleNamespace(returncode=-6, stdout=b"", stderr=b"compiler error")) + with pytest.raises(RuntimeError, match="compiler error"): + _native.invoke("cpu", "compile_cpu", b"input") + + +@pytest.mark.parametrize("code, exception", [(2, ValueError), (3, RuntimeError)]) +def test_worker_preserves_known_exception_types(monkeypatch, code, exception): + monkeypatch.setattr(default_tile_context.config, "compiler_timeout_sec", 1) + monkeypatch.setattr(_native, "compiler", lambda _: SimpleNamespace(__file__="/compiler")) + monkeypatch.setattr(_native.subprocess, "run", lambda *args, **kwargs: + SimpleNamespace(returncode=code, stdout=b"", stderr=b"native error")) + with pytest.raises(exception) as caught: + _native.invoke("cpu", "compile_cpu", b"input") + assert str(caught.value) == "native error" + + +def test_file_fingerprint_tracks_rebuilds_and_resolves_symlinks(tmp_path): + path = tmp_path / "compiler.so" + path.write_bytes(b"first") + alias = tmp_path / "alias.so" + alias.symlink_to(path) + first = _native.file_fingerprint(str(path)) + assert _native.file_fingerprint(str(alias)) == first + path.write_bytes(b"rebuilt compiler") + assert _native.file_fingerprint(str(path)) != first + + +@ct.kernel +def _empty(): + pass + + +def _divide_with_rounding(mode): + @ct.kernel + def divide(x, y, out): + a = ct.load(x, index=(0,), shape=(1,)) + b = ct.load(y, index=(0,), shape=(1,)) + ct.store(out, index=(0,), tile=ct.truediv(a, b, rounding_mode=mode)) + + return divide + + +@pytest.mark.parametrize("mode", [ct.RoundingMode.RZ, ct.RoundingMode.RM, + ct.RoundingMode.RP]) +@pytest.mark.parametrize("timeout", [None, 30]) +def test_native_rejects_unsupported_rounding(native, bytecode, monkeypatch, mode, timeout): + kernel = _divide_with_rounding(mode) + args = tuple(np.ones(1, dtype=np.float32) for _ in range(3)) + data, signature = bytecode(kernel, args, cpu) + monkeypatch.setattr(default_tile_context.config, "compiler_timeout_sec", timeout) + with pytest.raises(ValueError, match=f"unsupported rounding mode {mode.value}"): + cpu.compile(data, signature, cpu.Options(False)) + + +def test_native_worker_compiles_bytecode(native, bytecode, monkeypatch): + backend, options = cpu, cpu.Options(False) + data, signature = bytecode(_empty, (), backend) + monkeypatch.setattr(default_tile_context.config, "compiler_timeout_sec", 30) + assert backend.compile(data, signature, options).startswith(b"\x7fELF") diff --git a/backends/test/test_setup_editable.py b/backends/test/test_setup_editable.py new file mode 100644 index 00000000..73da81db --- /dev/null +++ b/backends/test/test_setup_editable.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import os +from pathlib import Path +import runpy + +import pytest +import setuptools + + +@pytest.fixture +def link_extension(monkeypatch): + monkeypatch.setattr(setuptools, "setup", lambda **kwargs: None) + setup = Path(__file__).resolve().parents[2] / "setup.py" + return runpy.run_path(str(setup))["_link_editable_extension"] + + +def test_existing_editable_link_is_reused(link_extension, tmp_path): + source = tmp_path / "source.so" + source.write_bytes(b"new") + destination = tmp_path / "extension.so" + destination.symlink_to(source) + + link_extension(str(source), str(destination)) + + assert destination.is_symlink() + assert os.readlink(destination) == str(source) + assert source.read_bytes() == b"new" + + +@pytest.mark.parametrize("old_destination", ["other-link", "broken-link", "regular-file"]) +def test_editable_link_replaces_stale_destination(link_extension, tmp_path, old_destination): + source = tmp_path / "source.so" + source.write_bytes(b"new") + destination = tmp_path / "extension.so" + if old_destination == "regular-file": + destination.write_bytes(b"old") + else: + old_source = tmp_path / "old.so" + if old_destination == "other-link": + old_source.write_bytes(b"old") + destination.symlink_to(old_source) + + link_extension(str(source), str(destination)) + + assert destination.is_symlink() + assert os.readlink(destination) == str(source) + assert destination.read_bytes() == b"new" + + +def test_missing_source_preserves_editable_link(link_extension, tmp_path): + destination = tmp_path / "extension.so" + destination.write_bytes(b"old") + + with pytest.raises(FileNotFoundError): + link_extension(str(tmp_path / "missing.so"), str(destination)) + + assert destination.read_bytes() == b"old" diff --git a/backends/test/test_signature.py b/backends/test/test_signature.py new file mode 100644 index 00000000..82ea5a56 --- /dev/null +++ b/backends/test/test_signature.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 +"""Keep the Python signature builder of custom backends in line with the CUDA launcher.""" + +from typing import Annotated + +import numpy as np +import pytest + +import cuda.tile as ct +from cuda.tile._backends._signature import build_signature, constant_parameters +from cuda.tile.compilation import ( + ArrayConstraint, CallingConvention, ConstantConstraint, KernelSignature, ScalarConstraint) + +_StaticStrides = Annotated[ct.Array, ct.ArrayAnnotation(static_stride_dims=(0,))] +_StaticShape = Annotated[ct.Array, ct.ArrayAnnotation(static_shape_dims=(-1,))] + + +@ct.kernel +def _kernel(x, y: ct.IndexedWithInt64, z: _StaticStrides, flag: bool, scale: float, + count: int, wide: ct.ScalarInt64, tile: ct.Constant[int]): + pass + + +@ct.kernel +def _static_shape_kernel(x: _StaticShape): + pass + + +@ct.kernel +def _tuple_kernel(pair: tuple[int, int]): + pass + + +def test_constant_parameters(): + assert constant_parameters(_kernel) == (False,) * 7 + (True,) + assert constant_parameters(_tuple_kernel) == (False,) + + +def test_tuple_parameters_are_cuda_only(): + with pytest.raises(TypeError, match="tuple parameters"): + build_signature(_tuple_kernel, ((1, 2),)) + + +def test_build_signature_follows_annotations(): + x = np.zeros((32, 16), dtype=np.float32) + signature = build_signature(_kernel, (x, x, x, True, 1.5, 3, 1 << 40, 16)) + plain, int64_indexed, static_strides, *scalars, tile = signature.parameters + + assert all(isinstance(p, ArrayConstraint) for p in (plain, int64_indexed, static_strides)) + assert (plain.index_dtype, int64_indexed.index_dtype) == (ct.int32, ct.int64) + assert plain.stride_constant == (None, 1) + # Annotated static strides replace the inference of unit strides. + assert static_strides.stride_constant == (16, None) + assert plain.shape_constant == (None, None) + assert scalars == [ScalarConstraint(ct.bool_), ScalarConstraint(ct.float32), + ScalarConstraint(ct.int32), ScalarConstraint(ct.int64)] + assert tile == ConstantConstraint(16) + + +def test_static_shape_dims_require_cutile_python_v2(): + x = np.zeros((32, 16), dtype=np.float32) + with pytest.raises(ValueError): + build_signature(_static_shape_kernel, (x,), + calling_convention=CallingConvention.cutile_python_v1()) + signature = build_signature(_static_shape_kernel, (x,)) + assert signature.calling_convention == CallingConvention.cutile_python_v2() + assert signature.parameters[0].shape_constant == (None, 16) + + +def test_size_one_dims_become_static_shapes_as_for_cuda_launches(): + x = np.zeros((32, 16), dtype=np.float32) + default = build_signature(_kernel, (x[:1], x, x, True, 1.5, 3, 1, 16)) + assert default.calling_convention == CallingConvention.cutile_python_v2() + assert default.parameters[0].shape_constant == (1, None) + assert default.parameters[1].shape_constant == (None, None) + + unit_free = build_signature(_kernel, (x, x, x, True, 1.5, 3, 1, 16)) + assert unit_free.calling_convention == CallingConvention.cutile_python_v1() + + v1 = build_signature(_kernel, (x[:1], x, x, True, 1.5, 3, 1, 16), + calling_convention=CallingConvention.cutile_python_v1()) + assert v1.parameters[0].shape_constant == (None, None) + + +@ct.kernel +def _cuda_kernel(x, y, flag: bool, scale: float, count: int, wide: ct.ScalarInt64, + big: ct.IndexedWithInt64, z: _StaticStrides, tile: ct.Constant[int]): + pass + + +def test_build_signature_matches_the_cuda_launcher(): + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("needs a CUDA device") + base = torch.zeros((64, 64), device="cuda") + for x, y in [(base, base[1:, ::2]), (base[:, :17], base.t()), + (base.as_strided((16, 16), (0, 1)), base.view(-1)[1:]), + (base[:1], base[:, :1])]: + args = (x, y, True, 1.5, 3, 1 << 40, y, x, 16) + # cutile_python_v2 enables the static size-one shapes of CUDA launches. + for convention in (CallingConvention.cutile_python_v1(), + CallingConvention.cutile_python_v2()): + assert build_signature(_cuda_kernel, args, calling_convention=convention) \ + == KernelSignature.from_kernel_args(_cuda_kernel, args, convention) diff --git a/changelog.d/autotune.md b/changelog.d/autotune.md new file mode 100644 index 00000000..685e9db4 --- /dev/null +++ b/changelog.d/autotune.md @@ -0,0 +1,6 @@ + + + +- Added `ct.autotune`, a decorator that tunes `ct.Constant` parameters and compiler hints of a + kernel, or the parameters of a factory preparing several launches, by benchmarking each + `ct.tune.Config` once per input layout and reusing the best one. diff --git a/changelog.d/custom-backend.md b/changelog.d/custom-backend.md new file mode 100644 index 00000000..fc3691cc --- /dev/null +++ b/changelog.d/custom-backend.md @@ -0,0 +1,10 @@ + + + +- Added an experimental CPU backend, built with `CUTILE_BUILD_VARIANT=cpu` + (see `backends/README.md`). `ct.set_backend("cpu")`, + or any module implementing the `Backend` protocol from + `cuda.tile._backends._custom`, runs `ct.launch` on that backend regardless of + argument devices; `ct.set_backend(None)` restores CUDA. + `ct.backend_options({...})` scopes backend options for compiles and launches; + unknown options are rejected. diff --git a/pyproject.toml b/pyproject.toml index af26df11..7fd87b93 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,8 @@ include = [ "cuda.tile._passes", "cuda.tile._bytecode", "cuda.tile._library_lifecycle", + "cuda.tile._backends", + "cuda.tile._backends.cpu", "cuda.tile.tune", "cuda.tile.jax", ] @@ -72,5 +74,55 @@ include = [ "cuda.tile" = ["VERSION", "_cext.pyi"] +# Development environments of the non-CUDA variants, which build cuda-tile from +# source without build isolation (see backends/README.md): +# CUTILE_BUILD_VARIANT=cpu uv sync --group cpu --no-build-isolation-package cuda-tile +# Dependency groups are not part of the package metadata. +[dependency-groups] +backend-build = [ + # Without build isolation, [build-system] requirements must be installed too. + "setuptools==84.0.0", + "wheel", + "cmake>=3.24,<4", + "ninja", + "nanobind>=2.9,<3", + "pybind11>=2.13.1", + # CUDA headers for building _cext without a CUDA Toolkit. + "cuda-toolkit[cudart,nvvm]>=13.2,<13.6; sys_platform == 'linux'", +] +cpu = [ + { include-group = "backend-build" }, + "numpy", + "torch", +] + [tool.uv] managed = true +# Without this, uv derives a key from `packages.find.where` and stores it as an +# mtime while comparing it as an inode, so it rebuilds on every sync. +cache-keys = [ + { file = "pyproject.toml" }, + { file = "setup.py" }, + { file = "CMakeLists.txt" }, + { file = "cmake/**" }, + { file = "cext/**" }, + # Native sources of the non-CUDA variants. + { file = "backends/**/CMakeLists.txt" }, + { file = "backends/**/*.cmake" }, + { file = "backends/**/*.in" }, + { file = "backends/**/*.cpp" }, + { file = "backends/**/*.h" }, + { env = "CUTILE_BUILD_VARIANT" }, + { env = "CMAKE_ARGS" }, + { env = "USE_LLVM_REVISION" }, +] + +[[tool.uv.index]] +name = "pytorch-cpu" +url = "https://download.pytorch.org/whl/cpu" +explicit = true + +[tool.uv.sources] +torch = [ + { index = "pytorch-cpu", group = "cpu" }, +] diff --git a/samples/cpu/AttentionFMHA-cpu.py b/samples/cpu/AttentionFMHA-cpu.py new file mode 100644 index 00000000..2c2cc25d --- /dev/null +++ b/samples/cpu/AttentionFMHA-cpu.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import os +import sys +import cuda.tile as ct +import torch +import math + +from torch.nn.functional import scaled_dot_product_attention + +# Kernels and helpers shared with the CUDA sample in the parent directory. +sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")) +from AttentionFMHA import fmha_kernel # noqa: E402 + + +ct.set_backend("cpu") + + +# --- Wrapper function to launch the FMHA kernel --- +def cutile_fmha(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, + qk_scale: float | None = None, + input_pos: int = 0, + tile_m: int = 128, + tile_n: int = 128, + query_group_size: int = 1, + causal: bool = False, + num_cpu_threads: int = 0, + kernel=fmha_kernel) -> torch.Tensor: + """ + Performs Fused Multi-Head Attention (FMHA) using a cuTile kernel. + + Args: + Q (torch.Tensor): Query tensor (Batch, Heads, SeqLen_Q, D_k). + K (torch.Tensor): Key tensor (Batch, KV_Heads, SeqLen_KV, D_k). + V (torch.Tensor): Value tensor (Batch, KV_Heads, SeqLen_KV, D_v). + qk_scale (float, optional): Scaling factor for QK dot product. Defaults to 1/sqrt(D_k). + input_pos (int, optional): Global start pos for queries (causal masking). Defaults to 0. + tile_m (int): Tile size for Query sequence length (M dimension). + tile_n (int): Tile size for Key/Value sequence length (N dimension). + query_group_size (int): Number of query heads per key/value head. + causal (bool): If True, applies causal masking. + num_cpu_threads (int): Number of CPU threads used to run the kernel. + + Returns: + torch.Tensor: Output tensor (Batch, Heads, SeqLen_Q, D_v). + """ + # --- Input Validation --- + if Q.ndim != 4 or K.ndim != 4 or V.ndim != 4: + raise ValueError("Input tensors Q, K, V must be 4D (Batch, Heads, SeqLen, Dim).") + if Q.shape[0] != K.shape[0] or Q.shape[0] != V.shape[0]: + raise ValueError("Batch dimensions must match for Q, K, V.") + if Q.shape[1] % query_group_size != 0: + raise ValueError("Number of query heads must be divisible by query_group_size.") + if K.shape[1] * query_group_size != Q.shape[1]: + raise ValueError("K_Heads * query_group_size must equal Q_Heads.") + if Q.shape[3] != K.shape[3]: + raise ValueError("D_k (last dim of Q and K) must match.") + if K.shape[2] != V.shape[2]: + raise ValueError("SeqLen_KV (dim 2 of K and V) must match.") + if Q.device != K.device or Q.device != V.device or not Q.is_cpu: + raise ValueError("All input tensors must be on the CPU.") + if Q.dtype != K.dtype or Q.dtype != V.dtype: + raise ValueError("All input tensors must have the same data type.") + + Batch, Heads, SeqLen_Q, D_k = Q.shape + _, KV_Heads, SeqLen_KV, D_v = V.shape + if D_v != D_k: + raise ValueError("The FMHA kernel requires matching query/key and value head dimensions.") + even_k = (SeqLen_KV % tile_n) == 0 + + if qk_scale is None: + qk_scale = 1.0 / math.sqrt(D_k) + + # --- Create Output Tensor --- + Out = torch.empty((Batch, Heads, SeqLen_Q, D_v), dtype=torch.float32, device=Q.device) + + # --- Calculate Grid Dimensions --- + grid_x = math.ceil(SeqLen_Q / tile_m) + grid_y = Batch * Heads + grid = (grid_x, grid_y, 1) + + # --- Launch the FMHA Kernel --- + with ct.backend_options({"num_cpu_threads": num_cpu_threads}): + ct.launch(None, grid, kernel, ( + Q, K, V, Out, + qk_scale, + input_pos, + D_k, + Heads, + tile_m, + tile_n, + query_group_size, + causal, + even_k + )) + + return Out + + +def torch_fmha(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, + is_causal: bool, enable_gqa: bool) -> torch.Tensor: + return scaled_dot_product_attention( + Q, K, V, is_causal=is_causal, enable_gqa=enable_gqa + ).to(torch.float32) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--correctness-check", + action="store_true", + help="Check the correctness of the results", + ) + parser.add_argument("--num-cpu-threads", type=int, default=0) + args = parser.parse_args() + print("--- Running cuTile Fused Multi-Head Attention (FMHA) Sample ---") + + # --- User Configuration --- + BATCH_SIZE = 2 + NUM_HEADS = 8 + SEQ_LEN_Q = 128 + SEQ_LEN_KV = 128 + D_K = 32 + D_V = 32 + + QUERY_GROUP_SIZE = 1 + + DTYPE = torch.bfloat16 + + Q_input = torch.randn(BATCH_SIZE, NUM_HEADS, SEQ_LEN_Q, D_K, dtype=DTYPE, device='cpu') + K_input = torch.randn(BATCH_SIZE, NUM_HEADS // QUERY_GROUP_SIZE, SEQ_LEN_KV, D_K, + dtype=DTYPE, device='cpu') + V_input = torch.randn(BATCH_SIZE, NUM_HEADS // QUERY_GROUP_SIZE, SEQ_LEN_KV, D_V, + dtype=DTYPE, device='cpu') + + print(" Configuration:") + print(f" Batch Size: {BATCH_SIZE}") + print(f" Number of Heads: {NUM_HEADS}") + print(f" Query Sequence Length: {SEQ_LEN_Q}") + print(f" KV Sequence Length: {SEQ_LEN_KV}") + print(f" Head Dimension (D_k): {D_K}") + print(f" Value Dimension (D_v): {D_V}") + print(f" Data Type: {DTYPE}") + print(f"Input Q shape: {Q_input.shape}") + print(f"Input K shape: {K_input.shape}") + print(f"Input V shape: {V_input.shape}") + + # Test 1: Non-Causal Attention + print("\n--- Test 1: Non-Causal Attention ---") + output_fmha_cutile_non_causal = cutile_fmha( + Q=Q_input, K=K_input, V=V_input, + tile_m=32, tile_n=32, # Increased tile sizes + causal=False, + query_group_size=QUERY_GROUP_SIZE, + num_cpu_threads=args.num_cpu_threads + ) + print(f"""cuTile FMHA Output shape (Non-Causal):{output_fmha_cutile_non_causal.shape}, + dtype:{output_fmha_cutile_non_causal.dtype}""") + if args.correctness_check: + ref_fmha = torch_fmha(Q_input, K_input, V_input, + is_causal=False, enable_gqa=False) + torch.testing.assert_close(output_fmha_cutile_non_causal, ref_fmha, atol=1e-2, rtol=1e-2) + print("Correctness check passed") + else: + print("Correctness check disabled") + + # Test 2: Causal Attention + print("\n--- Test 2: Causal Attention ---") + output_fmha_cutile_causal = cutile_fmha( + Q=Q_input, K=K_input, V=V_input, + tile_m=128, tile_n=128, # Increased tile sizes + causal=True, + query_group_size=QUERY_GROUP_SIZE, + num_cpu_threads=args.num_cpu_threads + ) + print(f"""cuTile FMHA Output shape (Causal): {output_fmha_cutile_causal.shape}, + dtype: {output_fmha_cutile_causal.dtype}""") + if args.correctness_check: + ref_fmha = torch_fmha(Q_input, K_input, V_input, + is_causal=True, enable_gqa=False) + torch.testing.assert_close(output_fmha_cutile_causal, ref_fmha, atol=1e-2, rtol=5e-2) + print("Correctness check passed") + else: + print("Correctness check disabled") + + print("\n--- cuTile Fused Multi-Head Attention (FMHA) Sample execution complete ---") diff --git a/samples/cpu/BatchMatMul-cpu.py b/samples/cpu/BatchMatMul-cpu.py new file mode 100644 index 00000000..05e5d2d9 --- /dev/null +++ b/samples/cpu/BatchMatMul-cpu.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import os +import sys +from math import ceil +import cuda.tile as ct +import torch + +# Kernels and helpers shared with the CUDA sample in the parent directory. +sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")) +from BatchMatMul import batch_matmul_kernel # noqa: E402 + + +ct.set_backend("cpu") + + +def bmm(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype, + num_cpu_threads: int = 0) -> torch.Tensor: + """ + Batch Matrix Multiplication using cuTile's standard tiled kernel. + + Args: + a (torch.Tensor): Input tensor A with shape (Batch, M, K). + b (torch.Tensor): Input tensor B with shape (Batch, K, N). + + Returns: + Output tensor C with shape (Batch, M, N). + """ + # --- Input Validation --- + if a.ndim != 3 or b.ndim != 3: + raise ValueError("Input tensors for BMM must be 3D (Batch, M, K) and (Batch, K, N).") + if a.shape[0] != b.shape[0]: + raise ValueError(f"""Batch dimensions must match: + A.shape[0]={a.shape[0]}, B.shape[0]={b.shape[0]}.""") + if a.device != b.device or not a.is_cpu or not b.is_cpu or a.dtype != b.dtype: + raise ValueError("""Input tensors must be on the CPU + and have the same data type.""") + + # Get M, K, N dimensions + Batch, M, K = a.shape + _, K_b, N = b.shape + assert K == K_b, f"Incompatible K dimensions: A's K is {K}, B's K is {K_b}" + + # Create output tensor + output = torch.empty((Batch, M, N), device=a.device, dtype=out_dtype) + + tm_val, tn_val, tk_val = ((32, 32, 32) if a.dtype == torch.bfloat16 + else (128, 256, 128)) + + # --- Grid calculation for standard 3D tiled kernel --- + grid = (Batch, ceil(M / tm_val), ceil(N / tn_val)) + + # --- Launch kernel --- + assume_in_bounds = (a.dtype == torch.bfloat16 and M % tm_val == 0 and + N % tn_val == 0 and K % tk_val == 0) + with ct.backend_options({"num_cpu_threads": num_cpu_threads, + "assume_in_bounds": assume_in_bounds}): + ct.launch(None, grid, batch_matmul_kernel, + (a, b, output, tm_val, tn_val, tk_val)) + + return output + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--correctness-check", + action="store_true", + help="Check the correctness of the results", + ) + parser.add_argument("--num-cpu-threads", type=int, default=0) + args = parser.parse_args() + print("--- Running cuTile Batched Matrix Multiplication (Standard Tiled) Sample ---") + + # --- User Configuration for BMM Example --- + BATCH_DIM = 4 + M_DIM = 512 + K_DIM = 256 + N_DIM = 1024 + + A = torch.randn(BATCH_DIM, M_DIM, K_DIM, dtype=torch.bfloat16, device='cpu') + B = torch.randn(BATCH_DIM, K_DIM, N_DIM, dtype=torch.bfloat16, device='cpu') + print(f"Input A shape: {A.shape}, dtype: {A.dtype}") + print(f"Input B shape: {B.shape}, dtype: {B.dtype}") + + C_bmm_cutile = bmm(A, B, A.dtype, args.num_cpu_threads) + print(f"""cuTile Standard BMM Output C + shape:{C_bmm_cutile.shape}, + dtype: {C_bmm_cutile.dtype}""") + if args.correctness_check: + torch.testing.assert_close(C_bmm_cutile, A @ B, atol=1e-1, rtol=1e-2) + print("Correctness check passed") + else: + print("Correctness check disabled") + + print("\n--- cuTile Batched Matrix Multiplication (Standard Tiled) examples complete ---") diff --git a/samples/cpu/BlockScaledMatMul-cpu.py b/samples/cpu/BlockScaledMatMul-cpu.py new file mode 100644 index 00000000..5e0483a5 --- /dev/null +++ b/samples/cpu/BlockScaledMatMul-cpu.py @@ -0,0 +1,325 @@ +# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import os +import sys +import cuda.tile as ct +import torch + +# Kernels and helpers shared with the CUDA sample in the parent directory. +sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")) +from BlockScaledMatMul import swizzle_32_4_4, swizzle_2d # noqa: E402 + + +ct.set_backend("cpu") + + +def block_quantize(x: torch.Tensor, block_size: int, dtype: torch.dtype = torch.float16): + + """ + Args: + x (torch.Tensor): input tensor. + block_size (int): size of block. + dtype (torch.dtype): datatype used to carry quantized values and scales on CPU. + + Returns: + tuple[torch.Tensor, torch.Tensor, int]: A tuple containing: + - x (torch.Tensor): in dtype. + - scale (torch.Tensor): in dtype. + + Raises: + ValueError: If the requested block size is larger than the inner most + dimension of the input tensor. + """ + + dtype_max = torch.finfo(torch.float8_e4m3fn).max + + if x.shape[-1] < block_size: + raise ValueError(f'The requested block size {block_size} is larger than the inner most ' + f'dimension of the input tensor {x.shape[-1]}') + + if x.shape[-1] % block_size != 0: + print('[WARNING] block size is not a multiple of the inner most ' + 'dimension of the input tensor') + print(' padding the input tensor...') + pad_len = (block_size - (x.shape[-1] % block_size)) % block_size + x = torch.nn.functional.pad(x, (0, pad_len), value=0) + + x_block = x.reshape(*x.shape[:-1], x.shape[-1] // block_size, block_size) + + scale = torch.max(x_block.abs(), dim=-1, keepdims=True)[0] + scale = torch.clamp(scale / dtype_max, min=1e-12) + scale = torch.pow(2.0, torch.ceil(torch.log2(scale))) + + x_q = x_block / scale + + x_q = x_q.to(dtype).reshape(x.shape) + scale = scale.to(dtype).squeeze(-1) + + return x_q, scale + + +ConstInt = ct.Constant[int] + + +def unswizzle_32_4_4(tile_swizzled_scale): + ''' + Kernel-side inverse of ``swizzle_32_4_4``: take a tile loaded + from the host swizzled scale tensor and recover the ``(M, K_s)`` + logical scale view used by the block-scaled matmul. + ''' + m1, m2, k1 = 32, 4, 4 + m0, k0, _, _ = tile_swizzled_scale.shape + + return (tile_swizzled_scale.reshape((m0, k0, m1, m2, k1)) + .permute((0, 3, 2, 1, 4)) + .reshape((m0 * m2 * m1, k0 * k1))) + + +@ct.kernel +def block_scaled_matmul_kernel( + A, A_scale, B, B_scale, C, + tm: ConstInt, # Tile size along M dimension (rows of C) + tn: ConstInt, # Tile size along N dimension (columns of C) + tk: ConstInt, # Tile size along K dimension (inner product dimension) + scaling_block_size: ConstInt): + + """ + cuTile kernel for block-scaled matrix multiplication. + + Computes C = (A * A_scale) @ (B * B_scale), accumulating in float32. + Each TileBlock computes one tm x tn output tile. The K dimension is processed + in chunks of tk, with tks scale values per K tile. + + If packed swizzle scales are passed, they get unswizzled into logical 2D scale tiles. + + Args: + A: Input matrix A (M x K). + A_scale: 2D scale of (M, K // scaling_block_size) or swizzle scale of + (M // 32 // 4, K // scaling_block_size // 4, 32, 4, 4) reshaped + into (M // 32 // 4, K // scaling_block_size // 4, 32, 16). + B: Input matrix B (K x N). + B_scale: 2D scale of (K // scaling_block_size, N) or swizzled scale of + (N // 32 // 4, K // scaling_block_size // 4, 32, 4, 4) reshaped + into (N // 32 // 4, K // scaling_block_size // 4, 32, 16). + C: Output matrix C (M x N). + tm (ConstInt): The height of the output tile computed by this block. + Corresponds to rows of A and C. + tn (ConstInt): The width of the output tile computed by this block. + Corresponds to columns of B and C. + tk (ConstInt): The depth of the inner loop (K-dimension) tile size. + Corresponds to columns of A and rows of B. + scaling_block_size (ConstInt): the scaling block size. + """ + GROUP_SIZE_M = 8 + M = A.shape[0] + N = B.shape[1] + bidx, bidy = swizzle_2d(M, N, tm, tn, GROUP_SIZE_M) + + tks = tk // scaling_block_size + + # Calculate the total number of tiles along the K-dimension that need to be processed. + # `ct.num_tiles(A, axis=1, shape=(tm, tk))` means: + # "View A as an MxK tensor tiled by (tm, tk), and return the number of tiles along + # axis 1 (the K dimension)." + # We pass shape=(tm, tk) to describe the 2D tiling, only `tk` matters for axis=1. + num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk)) + + # Initialize an accumulator for the current output tile (tm x tn). + # It's common practice to use `float32` for accumulation even with `float16` inputs + # to maintain higher precision during the sum-reduction of the matrix multiplication. + accumulator = ct.full((tm, tn), 0, dtype=ct.float32) + zero_pad = ct.PaddingMode.ZERO + + # K-dimension loop: Iterate over the K-dimension in chunks of 'tk'. + # In each iteration, a `tm` x `tk` tile from A and a `tk` x `tn` tile from B + # are loaded, multiplied, and accumulated. + for k in range(num_tiles_k): + # Load tile from matrix A. + # The `index=(bidx, k_tile_idx)` specifies which (M-tile, K-tile) to load + # from global memory A. `shape=(tm, tk)` defines the size of this tile. + a = ct.load(A, index=(bidx, k), shape=(tm, tk), padding_mode=zero_pad) + + if len(A_scale.shape) == 2: + # 2D scale path. A_scale is already stored in logical shape (M, K_s). + a_scale = ct.load(A_scale, index=(bidx, k), shape=(tm, tks), padding_mode=zero_pad) + else: + # Load the packed scale tile, unswizzle it, and reshape to + # the logical block-scaled matmul shape (tm, tks). + + # unswizzle + a_scale_swizzled = ct.load(A_scale, index=(bidx, k, 0, 0), + shape=(tm // 32 // 4, tks // 4, 32, 16), + padding_mode=zero_pad) + a_scale = unswizzle_32_4_4(a_scale_swizzled) + + # Load tile from matrix B. + # The `index=(k_tile_idx, bidy)` specifies which (K-tile, N-tile) to load + # from global memory B. `shape=(tk, tn)` defines the size of this tile. + b = ct.load(B, index=(k, bidy), shape=(tk, tn), padding_mode=zero_pad) + + if len(B_scale.shape) == 2: + b_scale = ct.load(B_scale, index=(k, bidy), shape=(tks, tn), padding_mode=zero_pad) + else: + # B scales are stored N-major. Unswizzle it, reshape it to + # (tn, tks), then transpose it to the logical block-scaled shape (tks, tn). + + # unswizzle + b_scale_swizzled = ct.load(B_scale, index=(bidy, k, 0, 0), + shape=(tn // 32 // 4, tks // 4, 32, 16), + padding_mode=zero_pad) + b_scale = unswizzle_32_4_4(b_scale_swizzled).permute((1, 0)) + + # Expand one scale factor per K block and apply it before ordinary MMA. + a_scale = ct.broadcast_to(a_scale.reshape((tm, tks, 1)), + (tm, tks, scaling_block_size)).reshape((tm, tk)) + b_scale = ct.broadcast_to(b_scale.reshape((tks, 1, tn)), + (tks, scaling_block_size, tn)).reshape((tk, tn)) + accumulator = ct.mma(a * a_scale, b * b_scale, accumulator) + + # Store the computed tile to the global memory of the output matrix C. + # The `(bidx, bidy)` directly corresponds to the tile's position in the 2D output matrix. + ct.store(C, index=(bidx, bidy), tile=accumulator) + + +def cutile_block_scaled_matmul(A: torch.Tensor, A_scale: torch.Tensor, + B: torch.Tensor, B_scale: torch.Tensor, + num_cpu_threads: int = 0) -> torch.Tensor: + + """ + Performs block-scaled matrix multiplication using a cuTile kernel. + + This wrapper function handles input validation, determines appropriate + tile sizes based on data type, calculates the necessary grid dimensions, + and launches the `block_scaled_matmul_kernel`. + + Args: + A (torch.Tensor): The first input matrix (M x K). Must be on the CPU. + B (torch.Tensor): The second input matrix (K x N). Must be on the CPU + and have its K dimension match A's K dimension. + A_scale (torch.Tensor): Either 2D scale with shape (M, K // scaling_block_size) or + swizzled scale of (M // 32 // 4, + K // scaling_block_size // 4, 32, 16). + B_scale (torch.Tensor): Either 2D scale with shape (K // scaling_block_size, N) or + swizzled scale of (N // 32 // 4, + K // scaling_block_size // 4, 32, 16). + + Returns: + torch.Tensor: The resulting matrix C (M x N) on the CPU. + + Raises: + ValueError: If matrices are incompatible (K dimensions don't match), + or if they are not on the CPU. + """ + # --- Input Validation --- + if A.shape[1] != B.shape[0]: + raise ValueError(f"Incompatible matrices: K dimension of A ({A.shape[1]}) " + f"must match K dimension of B ({B.shape[0]})") + if A.device != B.device or A.device != A_scale.device or A.device != B_scale.device: + raise ValueError("Input tensors must be on the same device.") + if not A.is_cpu or not A_scale.is_cpu or not B.is_cpu or not B_scale.is_cpu: + raise ValueError("Input tensors must be on the CPU.") + # Note: cuTile handles dtype compatibility within the kernel, + # but inputs should generally match. + + # The generic CPU lowering unrolls whole tiles, so keep the swizzle-compatible + # tiles small: tm and tn are multiples of 128 and tk spans 4 scale blocks. + tm, tn, tk, scaling_block_size = 128, 128, 128, 32 + + # --- Get Matrix Dimensions --- + m, _ = A.shape + _, n = B.shape + + # --- Calculate Grid Dimensions for Kernel Launch (1D Grid) --- + # The grid defines how many blocks will be launched. + # Each block computes one (tm x tn) output tile of matrix C. + # `ct.cdiv(total_dim, tile_dim)` ensures enough blocks are launched to cover + # the entire matrix, even if dimensions are not perfect multiples of tile sizes. + grid_x = ct.cdiv(m, tm) # Number of blocks needed along the M dimension (rows of C) + grid_y = ct.cdiv(n, tn) # Number of blocks needed along the N dimension (columns of C) + grid_size = grid_x * grid_y + + grid = (grid_size, 1, 1) + + # --- Create Output Tensor C --- + # The output tensor `C` is initialized with the correct dimensions (M x N), + # on the same device, and with the same data type as the input matrices. + C = torch.empty((m, n), device=A.device, dtype=torch.float32) + + # --- Launch the cuTile Kernel --- + # The `block_scaled_matmul_kernel` is launched with the calculated grid dimensions. + # `tm`, `tn`, and `tk` are passed as Constant integers to the kernel. + kernel = block_scaled_matmul_kernel + with ct.backend_options({"num_cpu_threads": num_cpu_threads}): + ct.launch(None, grid, kernel, ( + A, A_scale, B, B_scale, C, tm, tn, tk, scaling_block_size)) + + return C + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--correctness-check", + action="store_true", + help="Check the correctness of the results", + ) + parser.add_argument("--num-cpu-threads", type=int, default=0) + args = parser.parse_args() + + # --- Running cuTile Block Scaled Matrix Multiplication Examples --- + print("--- Running cuTile Block Scaled Matrix Multiplication Examples ---") + + # Define common matrix dimensions for the examples + M_dim = 512 + N_dim = 512 + K_dim = 768 + + scaling_block_size = 32 + KS_dim = K_dim // scaling_block_size + + print("\n--- Test Case: Block Scaled Matrix Multiplication with M = 512, N = 512, " + "K = 768, Scaling Block Size = 32 ---") + + A = torch.rand((M_dim, K_dim), device='cpu') + B = torch.rand((N_dim, K_dim), device='cpu') + + A, A_scale = block_quantize(A, scaling_block_size) + B, B_scale = block_quantize(B, scaling_block_size) + + B = B.T + B_scale = B_scale.T + + k = A.shape[-1] + ks = k // scaling_block_size + + A_s_swizzled = swizzle_32_4_4(A_scale) + B_s_swizzled = swizzle_32_4_4(B_scale.T.contiguous()) + + print(f"Input A shape: {A.shape}, dtype: {A.dtype}") + print(f"Input B shape: {B.shape}, dtype: {B.dtype}") + + atol, rtol = 1e-4, 1e-3 + + # Perform matrix multiplication using the cuTile wrapper function. + C_cutile = cutile_block_scaled_matmul( + A, A_scale, B, B_scale, args.num_cpu_threads) + C_cutile_swizzled = cutile_block_scaled_matmul( + A, A_s_swizzled, B, B_s_swizzled, args.num_cpu_threads) + print(f"cuTile Output C shape: {C_cutile.shape}, dtype: {C_cutile.dtype}") + + if args.correctness_check: + ref_A_scale = torch.repeat_interleave(A_scale, scaling_block_size, dim=1).to(torch.float32) + ref_B_scale = torch.repeat_interleave(B_scale, scaling_block_size, dim=0).to(torch.float32) + ref = (A.to(torch.float32) * ref_A_scale) @ (B.to(torch.float32) * ref_B_scale) + + torch.testing.assert_close(C_cutile, ref, atol=atol, rtol=rtol) + torch.testing.assert_close(C_cutile_swizzled, ref, atol=atol, rtol=rtol) + print("Correctness check passed") + else: + print("Correctness check disabled") + + print("\n--- cuTile block scaled matrix multiplication example completed. ---") diff --git a/samples/cpu/FFT-cpu.py b/samples/cpu/FFT-cpu.py new file mode 100644 index 00000000..4f17d1a9 --- /dev/null +++ b/samples/cpu/FFT-cpu.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import os +import sys +import cuda.tile as ct +import torch + +# Kernels and helpers shared with the CUDA sample in the parent directory. +sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")) +from FFT import fft_kernel, make_twiddles # noqa: E402 + + +ct.set_backend("cpu") + + +# --- Wrapper function to launch the fft_kernel --- +def cutile_fft( + x: torch.Tensor, + factors: tuple, # (F0, F1, F2) - factors of N + atom_packing_dim: int = 64, # The 'D' parameter for data packing/unpacking + num_cpu_threads: int = 0, +) -> torch.Tensor: + """ + Performs a Batched 1D Fast Fourier Transform (FFT) using a cuTile kernel + based on multi-dimensional factorization (similar to a Cooley-Tukey algorithm). + + This function prepares the input data, generates necessary pre-computed + matrices, launches the cuTile kernel, and unpacks the results. + + Args: + x (torch.Tensor): Input tensor of shape (Batch, N) containing complex64 numbers. + This tensor *must* be on the CPU. + N (the FFT size) must be factorable into factors[0]*factors[1]*factors[2]. + factors (tuple): A tuple (F0, F1, F2) representing the factors of N. + The product F0 * F1 * F2 must equal N. These factors define + the logical 3D shape for the FFT decomposition within the kernel. + atom_packing_dim (int): The dimension 'D' used for data packing/unpacking + in the kernel. This value affects memory access patterns. + The total number of real/imaginary elements (N*2) must be + divisible by this dimension. Default is 64. + + Returns: + torch.Tensor: Output tensor of shape (Batch, N) containing the FFT results. + The output data type will be torch.complex64. + + Raises: + ValueError: If input tensor dimensions, device, or data type are incorrect, + if the provided factors do not multiply to N, or if N*2 is not + divisible by atom_packing_dim. + """ + # --- Input Validation --- + if x.ndim != 2: + raise ValueError("Input tensor must be 2D (Batch, N).") + if not x.is_cpu: + raise ValueError("Input tensor must be on the CPU.") + if x.dtype != torch.complex64: + raise ValueError("Input tensor dtype must be torch.complex64.") + + BS = x.shape[0] # Extract Batch Size from the input tensor's shape. + N = x.shape[1] # Extract Total FFT size from the input tensor's shape. + + F0, F1, F2 = factors + # Validate that the provided factors correctly decompose the total FFT size N. + if F0 * F1 * F2 != N: + raise ValueError(f"Factors ({F0}*{F1}*{F2}={F0*F1*F2}) do not multiply to N={N}. " + f"Please provide factors that correctly decompose N.") + + # Determine the underlying floating-point precision (e.g., float32) for + # the real and imaginary parts of the complex numbers. + PRECISION_DTYPE = x.real.dtype + + # --- Prepare Input Data for Kernel (Split real/imag, pack) --- + # Convert the complex input tensor (BS, N) to a real tensor (BS, N, 2) + # where the last dimension explicitly separates real and imaginary parts. + x_ri = torch.view_as_real(x) + + # Reshape the real/imaginary tensor to the packed format (BS, N*2 // D, D) + # that the kernel expects for efficient memory access. + # This step assumes that the total number of real/imaginary elements (N*2) + # is perfectly divisible by the `atom_packing_dim` (D). + if (N * 2) % atom_packing_dim != 0: + raise ValueError(f"Total real/imag elements (N*2 = {N*2}) must be divisible by " + f"atom_packing_dim ({atom_packing_dim}) for kernel packing.") + x_packed_in = x_ri.reshape(BS, N * 2 // atom_packing_dim, atom_packing_dim).contiguous() + + # --- Generate W (Rotation) and T (Twiddle) Matrices --- + # These matrices are pre-computed mathematically based on the FFT decomposition. + # They are generated on the same device as the input tensor. + W0_gmem, W1_gmem, W2_gmem, T0_gmem, T1_gmem = make_twiddles(factors, PRECISION_DTYPE, x.device) + + # --- Create Output Tensor --- + # Initialize an empty tensor with the same shape and properties as the packed input. + # This tensor will store the results computed by the kernel. + y_packed_out = torch.empty_like(x_packed_in) + + # --- Calculate Grid Dimensions --- + # For this FFT kernel, one thread block is launched for each item in the batch. + # The grid is a 3-tuple (grid_x, grid_y, grid_z). + grid = (BS, 1, 1) + + # --- Launch the cuTile Kernel --- + # The `fft_kernel` is launched on the CPU with the calculated grid dimensions. + # All necessary input tensors (packed data, W and T matrices) and constant parameters + # (N, F0, F1, F2, BS, D) are passed to the kernel. + with ct.backend_options({"num_cpu_threads": num_cpu_threads}): + ct.launch(None, grid, fft_kernel, + (x_packed_in, y_packed_out, + W0_gmem, W1_gmem, W2_gmem, + T0_gmem, T1_gmem, + N, F0, F1, F2, BS, atom_packing_dim)) + + # --- Unpack Output from Kernel (Reshape, combine real/imag) --- + # Reshape the packed output tensor back to (BS, N, 2) to separate real/imaginary parts. + y_ri = y_packed_out.reshape(BS, N, 2) + # Convert the real/imaginary pair tensor back to a complex tensor (torch.complex64). + y_complex = torch.view_as_complex(y_ri) + + return y_complex + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--correctness-check", + action="store_true", + help="Check the correctness of the results", + ) + parser.add_argument("--num-cpu-threads", type=int, default=0) + args = parser.parse_args() + print("--- Running cuTile FFT Example ---") + + # --- User Configuration --- + BATCH_SIZE = 2 + # Total FFT size (N) must be factorable into factors[0] * factors[1] * factors[2]. + # For example, N = 1024 can be factored as (8, 16, 8). + FFT_SIZE = 8 # A smaller FFT size for demonstration purposes. + FFT_FACTORS = (2, 2, 2) # Factors for N=8 (2 * 2 * 2 = 8). + # The 'D' parameter for data packing/unpacking in the kernel. + # For N=8, N*2=16. ATOM_PACKING_DIM=2 is a valid divisor (16 % 2 == 0). + ATOM_PACKING_DIM = 2 + + # Data type for the real/imaginary components (e.g., float32 for complex64). + PRECISION_DTYPE = torch.float32 + + # --- Create Sample Input Data --- + # Generate a random input tensor of complex64 numbers, placed on the CPU. + # `torch.manual_seed(0)` ensures reproducibility of the random numbers. + torch.manual_seed(0) + input_data_complex = torch.randn(BATCH_SIZE, FFT_SIZE, dtype=torch.complex64, device='cpu') + + print(" Configuration:") + print(f" FFT Size (N): {FFT_SIZE}") + print(f" Batch Size: {BATCH_SIZE}") + print(f" FFT Factors (F0,F1,F2): {FFT_FACTORS}") + print(f" Atom Packing Dimension (D): {ATOM_PACKING_DIM}") + print(f"Input data shape: {input_data_complex.shape}, dtype: {input_data_complex.dtype}") + + # Perform FFT using the custom cuTile kernel. + output_fft_cutile = cutile_fft( + x=input_data_complex, + factors=FFT_FACTORS, + atom_packing_dim=ATOM_PACKING_DIM, + num_cpu_threads=args.num_cpu_threads, + ) + print( + f"""\ncuTile FFT Output shape: {output_fft_cutile.shape}, + dtype: {output_fft_cutile.dtype}""") + if args.correctness_check: + torch.testing.assert_close(output_fft_cutile, torch.fft.fft(input_data_complex, axis=-1)) + print("Correctness check passed") + else: + print("Correctness check disabled") + + print("\n--- cuTile FFT example execution complete ---") diff --git a/samples/cpu/LayerNorm-cpu.py b/samples/cpu/LayerNorm-cpu.py new file mode 100644 index 00000000..51b118c2 --- /dev/null +++ b/samples/cpu/LayerNorm-cpu.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import os +import sys +import math + +import torch +import cuda.tile as ct + +# Kernels and helpers shared with the CUDA sample in the parent directory. +sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")) +from LayerNorm import ( # noqa: E402 + layer_norm_fwd, + bwd_helper, + layer_norm_bwd_dwdb, + torch_layer_norm, +) + + +ct.set_backend("cpu") + + +ConstInt = ct.Constant[int] + + +@ct.kernel +def layer_norm_bwd_dx_partial_dwdb(DX, DY, DW, DB, X, W, Mean, Rstd, TILE_N: ConstInt): + """ + Backward pass part 1: computes dX and partial dW/dB. + Stores one partial gradient row per input row. + + Args: + DX: Output gradient with respect to X (M, N). + DY: Input gradient with respect to Y (M, N). + DW: Partial gradient with respect to W (M, N). + DB: Partial gradient with respect to B (M, N). + X: Input tensor (M, N). + W: Weight tensor (N,). + Mean: Mean tensor (M,). + Rstd: Reciprocal standard deviation tensor (M,). + TILE_N: Tile size along N dimension. + """ + bid_m = ct.bid(0) + num_tiles = ct.num_tiles(X, axis=1, shape=(1, TILE_N)) + N = X.shape[1] + mean = ct.load(Mean, index=(bid_m,), shape=(1,)) + rstd = ct.load(Rstd, index=(bid_m,), shape=(1,)) + + c1 = ct.full((1, TILE_N), 0, dtype=ct.float32) + c2 = ct.full((1, TILE_N), 0, dtype=ct.float32) + for j in range(num_tiles): + # Compute reduction terms for dX + _, xhat, wdy = bwd_helper(X, W, DY, bid_m, j, mean, rstd, TILE_N, N) + c1 += xhat * wdy + c2 += wdy + c1 = ct.sum(c1, axis=1) / N + c2 = ct.sum(c2, axis=1) / N + + for j in range(num_tiles): + # Compute dX and partial dW, dB + tdy, xhat, wdy = bwd_helper(X, W, DY, bid_m, j, mean, rstd, TILE_N, N) + tdx = (wdy - (xhat * c1 + c2)) * rstd + ct.store(DX, index=(bid_m, j), tile=tdx.astype(DX.dtype)) + + partial_dw = (tdy * xhat).astype(DW.dtype) + partial_db = tdy.astype(DB.dtype) + + ct.store(DW, index=(bid_m, j), tile=partial_dw) + ct.store(DB, index=(bid_m, j), tile=partial_db) + + +# --- cuTile LayerNorm Wrapper ------------------------------------------------------ + +class CuTileLayerNorm(torch.autograd.Function): + """ + A PyTorch Autograd Function wrapper for the cuTile LayerNorm kernel. + This class manages the forward and backward passes, bridging PyTorch tensors + with the cuTile kernel launches. + """ + + @staticmethod + def forward(ctx, input, weight, bias, eps, num_cpu_threads): + """ + Forward pass for LayerNorm. + + Args: + ctx: Context object to save tensors for backward pass. + input: Input tensor (*, ..., N). + weight: Scale parameter (N,). + bias: Shift parameter (N,). + eps: Epsilon for numerical stability. + + Returns: + Output tensor with the same shape as input. + """ + # Flatten input to (M, N) + x = input.reshape(-1, input.shape[-1]) + y = torch.empty_like(x) + M, _ = x.shape + + # Allocate temporary buffers for mean and reciprocal standard deviation + mean = torch.empty(M, dtype=torch.float32, device=x.device) + rstd = torch.empty(M, dtype=torch.float32, device=x.device) + + TILE_N = 1024 + # Launch the forward kernel with a 1D grid (M blocks) + with ct.backend_options({"num_cpu_threads": num_cpu_threads}): + ct.launch(None, (M,), layer_norm_fwd, + (x, weight, bias, y, mean, rstd, eps, TILE_N)) + + # Save tensors needed for the backward pass + ctx.save_for_backward(x, weight, bias, mean, rstd) + ctx.TILE_N = TILE_N + ctx.num_cpu_threads = num_cpu_threads + + return y.reshape(*input.shape) + + @staticmethod + def backward(ctx, grad_output): + """ + Backward pass for LayerNorm. + + Computes gradients for input, weight, and bias using two kernels: + 1. Computes dX and partial reductions for dW and dB. + 2. Performs the final reduction for dW and dB. + + Args: + ctx: Context object containing saved tensors. + grad_output: Gradient tensor of loss w.r.t. output (*, ..., N). + + Returns: + Gradients for input, weight, bias, and None for eps. + """ + x, weight, bias, mean, rstd = ctx.saved_tensors + TILE_N = ctx.TILE_N + num_cpu_threads = ctx.num_cpu_threads + M, N = x.shape + + # Flatten gradient output to (M, N) + dy = grad_output.reshape(-1, grad_output.shape[-1]) + dx = torch.empty_like(dy) + + # Allocate one partial gradient row per input row. + dw = torch.empty((M, N), dtype=torch.float32, device=weight.device) + db = torch.empty((M, N), dtype=torch.float32, device=bias.device) + + # Launch the first backward kernel to compute dX and partial dW/dB + with ct.backend_options({"num_cpu_threads": num_cpu_threads}): + ct.launch(None, (M,), layer_norm_bwd_dx_partial_dwdb, + (dx, dy, dw, db, x, weight, mean, rstd, TILE_N)) + + final_dw = torch.empty((N,), dtype=weight.dtype, device=weight.device) + final_db = torch.empty((N,), dtype=bias.dtype, device=bias.device) + TILE_M = 32 + + # Launch the second backward kernel to reduce partial dW/dB + with ct.backend_options({"num_cpu_threads": num_cpu_threads}): + ct.launch(None, (math.ceil(N / TILE_N),), layer_norm_bwd_dwdb, + (dw, db, final_dw, final_db, TILE_M, TILE_N)) + + return dx.reshape(*grad_output.shape), final_dw, final_db, None, None + + +def cutile_layer_norm(x, weight, bias, eps, num_cpu_threads=0): + return CuTileLayerNorm.apply(x, weight, bias, eps, num_cpu_threads) + + +# --- Torch Reference Implementation ----------------------------------------- + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--correctness-check", + action="store_true", + help="Check the correctness of the cuTile LayerNorm output against a torch reference.", + ) + parser.add_argument("--num-cpu-threads", type=int, default=0) + args = parser.parse_args() + + print("--- Running cuTile LayerNorm Forward/Backward Sample ---") + + shape = (1024, 2048) + dtype = torch.float32 + weight = torch.randn(shape[-1], dtype=dtype, device='cpu', requires_grad=True) + bias = torch.randn(shape[-1], dtype=dtype, device='cpu', requires_grad=True) + x = -2.3 + 0.5 * torch.randn(shape, dtype=dtype, device='cpu') + dy = 0.1 * torch.randn_like(x) + x.requires_grad_(True) + eps = 1e-5 + + print(f"Input shape: {shape}, dtype: {dtype}, eps: {eps}") + + atol, rtol = 1e-2, 1e-2 + + print("\n--- Executing cuTile LayerNorm Forward ---") + y = cutile_layer_norm(x, weight, bias, eps, args.num_cpu_threads)\ + + print("\n--- Executing cuTile LayerNorm Backward ---") + y.backward(dy, retain_graph=True) + dx, dw, db = [_.grad.clone() for _ in [x, weight, bias]] + x.grad, weight.grad, bias.grad = None, None, None + + if args.correctness_check: + print("\n--- Running correctness check against torch reference ---") + y_ref = torch_layer_norm(x, weight, bias, eps) + torch.testing.assert_close(y, y_ref, atol=atol, rtol=rtol) + + y_ref.backward(dy, retain_graph=True) + dx_ref, dw_ref, db_ref = [_.grad.clone() for _ in [x, weight, bias]] + torch.testing.assert_close(dx, dx_ref, atol=atol, rtol=rtol) + torch.testing.assert_close(dw, dw_ref, atol=atol, rtol=rtol) + torch.testing.assert_close(db, db_ref, atol=atol, rtol=rtol) + print("Correctness check passed") + else: + print("Correctness check disabled (use --correctness-check to enable)") + + print("\n--- cuTile LayerNorm Sample complete ---") diff --git a/samples/cpu/MatMul-cpu.py b/samples/cpu/MatMul-cpu.py new file mode 100644 index 00000000..983644b0 --- /dev/null +++ b/samples/cpu/MatMul-cpu.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import os +import sys +import cuda.tile as ct +import torch +from math import ceil # Required for host-side grid calculation + +# Kernels and helpers shared with the CUDA sample in the parent directory. +sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")) +from MatMul import matmul_kernel # noqa: E402 + + +ct.set_backend("cpu") + + +def cutile_matmul(A: torch.Tensor, B: torch.Tensor, + num_cpu_threads: int = 0) -> torch.Tensor: + """ + Performs matrix multiplication C = A @ B using a cuTile kernel with a 2D grid. + + This wrapper function handles input validation, determines appropriate + tile sizes based on data type, calculates the necessary grid dimensions, + and launches the `matmul_kernel`. + + Args: + A (torch.Tensor): The first input matrix (M x K). Must be on the CPU. + B (torch.Tensor): The second input matrix (K x N). Must be on the CPU + and have its K dimension match A's K dimension. + Returns: + torch.Tensor: The resulting matrix C (M x N) on the CPU. + + Raises: + ValueError: If matrices are incompatible (K dimensions don't match), + or if they are not on the CPU. + """ + # --- Input Validation --- + if A.shape[1] != B.shape[0]: + raise ValueError(f"Incompatible matrices: K dimension of A ({A.shape[1]}) " + f"must match K dimension of B ({B.shape[0]})") + if A.device != B.device: + raise ValueError("Input tensors must be on the same device.") + if not A.is_cpu or not B.is_cpu: + raise ValueError("Input tensors must be on the CPU.") + # Note: cuTile handles dtype compatibility within the kernel, but inputs should generally match. + + # BF16 uses AMX-friendly tiles; other dtypes retain their generic tile sizes. + if A.dtype == torch.bfloat16: + tm, tn, tk = 32, 32, 32 + elif A.dtype.itemsize == 2: + tm, tn, tk = 128, 256, 64 + else: + tm, tn, tk = 32, 32, 32 + + # --- Get Matrix Dimensions --- + m, k_a = A.shape # M = total rows of A (and C), K_A = total columns of A + k_b, n = B.shape # K_B = total rows of B, N = total columns of B (and C) + # Note: k_a and k_b must be equal due to validation. This is the 'K' dimension. + + # --- Calculate Grid Dimensions for Kernel Launch (1D Grid) --- + # The grid defines how many blocks will be launched. + # Each block computes one (tm x tn) output tile of matrix C. + # `ceil(total_dim / tile_dim)` ensures enough blocks are launched to cover + # the entire matrix, even if dimensions are not perfect multiples of tile sizes. + grid_x = ceil(m / tm) # Number of blocks needed along the M dimension (rows of C) + grid_y = ceil(n / tn) # Number of blocks needed along the N dimension (columns of C) + grid_size = grid_x * grid_y + grid = (grid_size, 1, 1) + + # --- Create Output Tensor C --- + # The output tensor `C` is initialized with the correct dimensions (M x N), + # on the same device, and with the same data type as the input matrices. + C = torch.empty((m, n), device=A.device, dtype=A.dtype) + + # --- Launch the cuTile Kernel --- + # The `matmul_kernel` is launched with the calculated grid dimensions. + # `tm`, `tn`, and `tk` are passed as Constant integers to the kernel. + assume_in_bounds = m % tm == 0 and n % tn == 0 and k_a % tk == 0 + with ct.backend_options({ + "num_cpu_threads": num_cpu_threads, + "assume_in_bounds": assume_in_bounds, + }): + ct.launch(None, grid, matmul_kernel, (A, B, C, tm, tn, tk)) + + return C + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--correctness-check", + action="store_true", + help="Check the correctness of the results", + ) + parser.add_argument("--num-cpu-threads", type=int, default=0) + args = parser.parse_args() + + for M_dim, N_dim, K_dim in ((512, 512, 768), (992, 512, 704)): + A = torch.randn(M_dim, K_dim, dtype=torch.bfloat16, device="cpu") + B = torch.randn(K_dim, N_dim, dtype=torch.bfloat16, device="cpu") + C_cutile = cutile_matmul(A, B, num_cpu_threads=args.num_cpu_threads) + print(f"Input shapes: {A.shape}, {B.shape}") + print(f"cuTile output shape: {C_cutile.shape}, dtype: {C_cutile.dtype}") + if args.correctness_check: + torch.testing.assert_close(C_cutile, A @ B, atol=1e-1, rtol=1e-2) + print("Correctness check passed") + else: + print("Correctness check disabled") diff --git a/samples/cpu/MoE-cpu.py b/samples/cpu/MoE-cpu.py new file mode 100644 index 00000000..b6bba473 --- /dev/null +++ b/samples/cpu/MoE-cpu.py @@ -0,0 +1,253 @@ +# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import os +import sys +import math + +import torch +import cuda.tile as ct + +# Kernels and helpers shared with the CUDA sample in the parent directory. +sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")) +from MoE import ( # noqa: E402 + fused_moe_kernel, + silu_and_mul_kernel, + moe_align_tile_size_torch, + torch_moe, + next_power_of_2, +) + + +ct.set_backend("cpu") + + +# --- cuTile MoE Wrapper ------------------------------------------------------ +def cutile_moe( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + tile_m: int, + tile_n: int, + tile_k: int, + num_cpu_threads: int = 0, +) -> torch.Tensor: + """ + Executes a Mixture-of-Experts (MoE) forward pass using the fused cuTile kernel. + + Args: + hidden_states: Token activations, shape (num_tokens, hidden_size) + w1: Expert gate+up projection weights, + shape (num_experts, intermediate_size * 2, hidden_size) + w2: Expert down projection weights, + shape (num_experts, hidden_size, intermediate_size) + topk_weights: Router weights per token, shape (num_tokens, topk) + topk_ids: Expert indices per token, shape (num_tokens, topk) + tile_m/n/k: Tile sizes for cuTile kernel launch + num_cpu_threads: Number of CPU threads used to run each kernel + + Returns: + Tensor with the same shape/dtype as `hidden_states`. + """ + out_dtype = hidden_states.dtype + device = hidden_states.device + + num_tokens, hidden_size = hidden_states.shape + num_experts, _, intermediate_size = w2.shape + _, topk = topk_ids.shape + + if w1.shape[1] != intermediate_size * 2: + raise ValueError("w1 must have 2 * intermediate_size rows (gate + up projection)") + + intermediate_cache1 = torch.zeros( + (num_tokens, topk, intermediate_size * 2), + device=device, + dtype=out_dtype, + ) + intermediate_cache2 = torch.zeros( + (num_tokens * topk, intermediate_size), + device=device, + dtype=out_dtype, + ) + intermediate_cache3 = torch.zeros( + (num_tokens, topk, hidden_size), + device=device, + dtype=out_dtype, + ) + + sorted_token_ids, sorted_expert_ids = moe_align_tile_size_torch( + topk_ids, + tile_m, + num_experts, + ) + + invoke_fused_moe_kernel( + hidden_states, + w1, + intermediate_cache1, + topk_weights, + sorted_token_ids, + sorted_expert_ids, + mul_routed_weight=False, + num_token_replicas=topk, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + num_cpu_threads=num_cpu_threads, + ) + + invoke_silu_and_mul_kernel( + intermediate_cache1.view(-1, intermediate_cache1.shape[-1]), + intermediate_cache2, + num_cpu_threads=num_cpu_threads, + ) + + invoke_fused_moe_kernel( + intermediate_cache2, + w2, + intermediate_cache3, + topk_weights, + sorted_token_ids, + sorted_expert_ids, + mul_routed_weight=True, + num_token_replicas=1, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + num_cpu_threads=num_cpu_threads, + ) + + return torch.sum(intermediate_cache3, dim=1) + + +# --- Helper Utilities ------------------------------------------------------- +def invoke_fused_moe_kernel( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + topk_weights: torch.Tensor, + sorted_token_ids: torch.Tensor, + sorted_expert_ids: torch.Tensor, + mul_routed_weight: bool, + num_token_replicas: int, + tile_m: int, + tile_n: int, + tile_k: int, + num_cpu_threads: int = 0, +) -> None: + m = sorted_token_ids.shape[0] + n = B.shape[1] + + grid = (math.ceil(m / tile_m) * math.ceil(n / tile_n),) + topk_weights = topk_weights.view(-1) + C = C.view(-1, C.shape[2]) + + with ct.backend_options({"num_cpu_threads": num_cpu_threads}): + ct.launch( + None, + grid, + fused_moe_kernel, + ( + A, + B, + C, + topk_weights, + sorted_token_ids, + sorted_expert_ids, + num_token_replicas, + mul_routed_weight, + tile_m, + tile_n, + tile_k, + ), + ) + + +def invoke_silu_and_mul_kernel( + AB: torch.Tensor, + C: torch.Tensor, + num_cpu_threads: int = 0, +): + A, B = AB.chunk(2, dim=-1) + with ct.backend_options({"num_cpu_threads": num_cpu_threads}): + ct.launch( + None, + (AB.shape[0],), + silu_and_mul_kernel, + ( + A, + B, + C, + next_power_of_2(C.shape[-1]) + ) + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--correctness-check", + action="store_true", + help="Check the correctness of the cuTile MoE output against a torch reference.", + ) + parser.add_argument("--num-cpu-threads", type=int, default=0) + args = parser.parse_args() + + print("--- Running cuTile Mixture-of-Experts (MoE) Sample ---") + + num_tokens = 48 + hidden_size = 512 + num_experts = 64 + intermediate_size = 1024 + topk = 8 + dtype = torch.bfloat16 + + device = "cpu" + print( + f"Tokens: {num_tokens}, Hidden: {hidden_size}, " + f"Experts: {num_experts}, Intermediate: {intermediate_size}, " + f"TopK: {topk}, dtype: {dtype}" + ) + + hidden_states = torch.empty( + num_tokens, hidden_size, device=device, dtype=dtype + ).normal_(0, 0.5) + w1 = torch.empty( + num_experts, intermediate_size * 2, hidden_size, device=device, dtype=dtype + ).normal_(0, 0.1) + w2 = torch.empty( + num_experts, hidden_size, intermediate_size, device=device, dtype=dtype + ).normal_(0, 0.1) + + # Unique expert IDs for each token (no repeating elements per row) + topk_ids = torch.stack([ + torch.randperm(num_experts, device=device)[:topk] + for _ in range(num_tokens) + ]) + topk_weights = torch.softmax( + torch.randn(num_tokens, topk, device=device), dim=-1 + ).to(dtype) + + print("\n--- Executing cuTile MoE ---") + # The generic CPU lowering unrolls whole tiles (including the A gather), so + # small tiles keep compilation fast. + output_cutile = cutile_moe(hidden_states, w1, w2, topk_weights, topk_ids, + tile_m=32, tile_n=64, tile_k=32, + num_cpu_threads=args.num_cpu_threads) + print(f"cuTile MoE output shape: {output_cutile.shape}, " + "dtype: {output_cutile.dtype}") + + if args.correctness_check: + print("\n--- Running correctness check against torch reference ---") + ref_output = torch_moe(hidden_states, w1, w2, topk_weights, topk_ids) + torch.testing.assert_close(output_cutile, ref_output, rtol=1e-1, atol=1e-1) + print("Correctness check passed") + else: + print("Correctness check disabled (use --correctness-check to enable)") + + print("\n--- cuTile Mixture-of-Experts (MoE) Sample complete ---") diff --git a/samples/cpu/Transpose-cpu.py b/samples/cpu/Transpose-cpu.py new file mode 100644 index 00000000..26ce5faf --- /dev/null +++ b/samples/cpu/Transpose-cpu.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import os +import sys +import cuda.tile as ct +import torch +from math import ceil # Required for host-side grid calculation + +# Kernels and helpers shared with the CUDA sample in the parent directory. +sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")) +from Transpose import transpose_kernel # noqa: E402 + + +ct.set_backend("cpu") + + +def cutile_transpose(x: torch.Tensor, num_cpu_threads: int = 0) -> torch.Tensor: + """ + Performs matrix transposition C = X.T using a cuTile kernel. + + This wrapper function handles input validation, determines appropriate + tile sizes based on data type, calculates the necessary grid dimensions, + and launches the `transpose_kernel`. + + Args: + x (torch.Tensor): The input matrix (M x N). + This tensor *must* be 2D and on the CPU. + + Returns: + torch.Tensor: The transposed matrix (N x M) on the CPU. + + Raises: + ValueError: If the input tensor is not on the CPU or is not 2D. + """ + # --- Input Validation --- + if not x.is_cpu: + raise ValueError("Input tensor must be on the CPU.") + if x.ndim != 2: + raise ValueError("Transpose kernel currently supports only 2D tensors.") + + # --- Get Matrix Dimensions --- + m, n = x.shape # Original dimensions of the input matrix: M rows, N columns + + # --- Determine Tile Shapes for Optimization --- + # These are the tile sizes (`tm`, `tn`) for the *input* matrix `x`. + # cuTile allows for specifying tile sizes as compile-time constants. + # For half-precision data types (like `torch.float16` or `torch.bfloat16`, + # where `itemsize` is 2 bytes), larger tile sizes (e.g., 128x128) might be + # chosen to leverage Tensor Cores for improved performance. + # For full-precision (e.g., `torch.float32`, `itemsize` is 4 bytes), + # smaller, more general tile sizes (e.g., 32x32) are often used. + if x.dtype.itemsize == 2: # Likely torch.float16 or torch.bfloat16 + tm, tn = 128, 128 # Example optimal tile sizes for 16-bit data + else: # Likely torch.float32 or other + tm, tn = 32, 32 # Example optimal tile sizes for 32-bit data + + # --- Calculate Grid Dimensions for Kernel Launch (2D Grid) --- + # The grid defines how many blocks will be launched. + # Each block computes one `(tn x tm)` output tile, which corresponds to + # processing one `(tm x tn)` input tile. + # `grid_x`: Number of blocks needed along the M dimension (rows of original x). + # `grid_y`: Number of blocks needed along the N dimension (columns of original x). + # `ceil(dimension / tile_size)` ensures that enough blocks are launched to cover + # the entire matrix, even if dimensions are not perfect multiples of tile sizes. + grid_x = ceil(m / tm) + grid_y = ceil(n / tn) + grid = (grid_x, grid_y, 1) # cuTile expects a 3-tuple for grid dimensions (x, y, z) + + # --- Create Output Tensor y --- + # The output tensor `y` will have transposed dimensions (N x M) + # and the same device and data type as the input `x`. + y = torch.empty((n, m), device=x.device, dtype=x.dtype) + + # --- Launch the cuTile Kernel --- + # The `transpose_kernel` is launched on the CPU with the calculated grid dimensions. + # `tm` and `tn` are passed as Constant integer parameters to the kernel. + with ct.backend_options({"num_cpu_threads": num_cpu_threads}): + ct.launch(None, grid, transpose_kernel, (x, y, tm, tn)) + + return y + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--correctness-check", + action="store_true", + help="Check the correctness of the results", + ) + parser.add_argument("--num-cpu-threads", type=int, default=0) + args = parser.parse_args() + + print("--- Running cuTile Matrix Transposition Examples ---") + + # Define common matrix dimensions for the examples + M_dim = 1024 + N_dim = 512 + + # --- Test Case 1: float16 (Half-Precision) --- + print("\n--- Test Case 1: Matrix Transposition with float16 (Half-Precision) ---") + # Create a random input matrix with float16 data type on the CPU. + x_fp16 = torch.randn(M_dim, N_dim, dtype=torch.float16, device='cpu') + print(f"Input x shape: {x_fp16.shape}, dtype: {x_fp16.dtype}") + + # Perform transposition using the cuTile wrapper function. + y_fp16_cutile = cutile_transpose(x_fp16, args.num_cpu_threads) + print(f"cuTile Output y shape: {y_fp16_cutile.shape}, dtype: {y_fp16_cutile.dtype}") + if args.correctness_check: + torch.testing.assert_close(y_fp16_cutile, x_fp16.T) + print("Correctness check passed") + else: + print("Correctness check disabled") + + # --- Test Case 2: float32 (Single-Precision) --- + print("\n--- Test Case 2: Matrix Transposition with float32 (Single-Precision) ---") + # Create a random input matrix with float32 data type on the CPU. + x_fp32 = torch.randn(M_dim, N_dim, dtype=torch.float32, device='cpu') + print(f"Input x shape: {x_fp32.shape}, dtype: {x_fp32.dtype}") + + # Perform transposition using the cuTile wrapper function. + y_fp32_cutile = cutile_transpose(x_fp32, args.num_cpu_threads) + print(f"cuTile Output y shape: {y_fp32_cutile.shape}, dtype: {y_fp32_cutile.dtype}") + if args.correctness_check: + torch.testing.assert_close(y_fp32_cutile, x_fp32.T) + print("Correctness check passed") + else: + print("Correctness check disabled") + + # --- Test Case 3: Non-square matrix with non-multiple dimensions --- + print("\n--- Test Case 3: Matrix Transposition with Non-Square, Non-Multiple Dimensions ---") + # Define matrix dimensions that are not exact multiples of the default tile sizes (32x32). + # Demonstration that the `ceil` function in grid calculation correctly handles partial tiles. + M_dim_non_mult = 1000 + N_dim_non_mult = 500 + x_non_mult = torch.randn(M_dim_non_mult, N_dim_non_mult, dtype=torch.float32, device='cpu') + print(f"Input x shape: {x_non_mult.shape}, dtype: {x_non_mult.dtype}") + + y_non_mult_cutile = cutile_transpose(x_non_mult, args.num_cpu_threads) + print(f"cuTile Output y shape: {y_non_mult_cutile.shape}, dtype: {y_non_mult_cutile.dtype}") + if args.correctness_check: + torch.testing.assert_close(y_non_mult_cutile, x_non_mult.T) + print("Correctness check passed") + else: + print("Correctness check disabled") + + print("\n--- All cuTile matrix transposition examples completed. ---") diff --git a/samples/cpu/VectorAddition-cpu.py b/samples/cpu/VectorAddition-cpu.py new file mode 100644 index 00000000..26d33279 --- /dev/null +++ b/samples/cpu/VectorAddition-cpu.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import os +import sys +import cuda.tile as ct +import torch +import math + +# Kernels and helpers shared with the CUDA sample in the parent directory. +sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")) +from VectorAddition import ( # noqa: E402 + vec_add_kernel_1d, + vec_add_kernel_2d, + vec_add_kernel_1d_gather, + vec_add_kernel_2d_gather, +) + + +ct.set_backend("cpu") + + +# --- Wrapper Function to Dispatch to Kernels --- +def vec_add(a: torch.Tensor, b: torch.Tensor, use_gather: bool = False, + num_cpu_threads: int = 0) -> torch.Tensor: + """ + Performs element-wise addition of two tensors (vector or matrix) using + different cuTile kernels based on dimensionality and gather/scatter preference. + + This function acts as a high-level interface, handling input validation, + determining appropriate tile and grid dimensions, and dispatching to + the correct cuTile kernel. + + Args: + a (torch.Tensor): The first input tensor. Must be 1D or 2D and on the CPU. + b (torch.Tensor): The second input tensor. Must match 'a' in shape, device, and dtype. + use_gather (bool): If True, uses kernels with explicit gather/scatter and masking + for robust boundary handling (recommended for non-power-of-2 + or non-tile-aligned dimensions). + If False, uses direct tiled loads/stores, which are simpler + but assume dimensions are perfectly divisible by tile sizes. + num_cpu_threads (int): Number of CPU threads used to run each kernel. + + Returns: + torch.Tensor: The resulting tensor after element-wise addition. + + Raises: + ValueError: If input tensors have mismatched shapes, incorrect dimensions, + are not on the CPU, or have different data types. + """ + # --- Input Validation --- + if a.shape != b.shape: + raise ValueError("Input tensors must have the same shape.") + if a.dim() > 2 or b.dim() > 2: + raise ValueError("This function currently supports only 1D or 2D tensors.") + if a.device != b.device: + raise ValueError("Input tensors must be on the same device.") + if not a.is_cpu or not b.is_cpu: + raise ValueError("Input tensors must be on the CPU.") + if a.dtype != b.dtype: + raise ValueError("Input tensors must have the same data type.") + + # Create an empty output tensor on the same device and with the same dtype as inputs. + c = torch.empty_like(a) + + # --- Dispatch based on Dimensionality --- + if a.dim() == 1: + N = a.shape[0] # Get the total size of the 1D vector + + # Heuristic for TILE size: + # Choose a power of 2, up to 1024, that is greater than or equal to N. + # This helps in efficient memory access patterns on the GPU. + # Handle N=0 gracefully to avoid log2(0) errors. + TILE = min(1024, 2 ** math.ceil(math.log2(N))) if N > 0 else 1 + + # Calculate the grid dimensions for launching the kernel. + # `math.ceil(N / TILE)` determines the number of blocks needed to cover + # the entire vector. Each block processes a `TILE`-sized chunk. + grid = (math.ceil(N / TILE), 1, 1) # (blocks_x, blocks_y, blocks_z) + + kernel = vec_add_kernel_1d_gather if use_gather else vec_add_kernel_1d + with ct.backend_options({"num_cpu_threads": num_cpu_threads}): + ct.launch(None, grid, kernel, (a, b, c, TILE)) + else: # a.dim() == 2 (Matrix) + M, N = a.shape # Get rows (M) and columns (N) of the matrix + + # Heuristic for 2D TILE sizes: + # TILE_Y is chosen as a power of 2, up to 1024, based on the column dimension. + TILE_Y = min(1024, 2 ** math.ceil(math.log2(N))) if N > 0 else 1 + # TILE_X is derived to keep the total elements processed by a block + # (TILE_X * TILE_Y) around 1024 (a common block size limit for threads). + TILE_X = max(1, 1024 // TILE_Y) + + # Calculate the 2D grid dimensions for launching the kernel. + # `math.ceil(M / TILE_X)` blocks along rows, `math.ceil(N / TILE_Y)` blocks along columns. + grid = (math.ceil(M / TILE_X), math.ceil(N / TILE_Y), 1) + + kernel = vec_add_kernel_2d_gather if use_gather else vec_add_kernel_2d + with ct.backend_options({"num_cpu_threads": num_cpu_threads}): + ct.launch(None, grid, kernel, (a, b, c, TILE_X, TILE_Y)) + + return c # Return the computed output tensor + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--correctness-check", + action="store_true", + help="Check the correctness of the results", + ) + parser.add_argument("--num-cpu-threads", type=int, default=0) + args = parser.parse_args() + + print("--- Running cuTile Vector/Matrix Addition Examples ---") + + # --- User Configuration --- + VECTOR_SIZE_1D = 1_000_000 + MATRIX_SHAPE_2D = (2048, 1024) # Rows, Columns + + # --- Test Case 1: 1D Vector Add (Direct Tiled) --- + print("\n--- Test 1: 1D Vector Add (Direct Tiled) ---") + # Create random input tensors on the CPU. + a_1d_direct = torch.randn(VECTOR_SIZE_1D, dtype=torch.float32, device='cpu') + b_1d_direct = torch.randn(VECTOR_SIZE_1D, dtype=torch.float32, device='cpu') + print(f"Input 1D shape: {a_1d_direct.shape}, dtype: {a_1d_direct.dtype}") + + # Call the vec_add wrapper function, requesting the direct tiled kernel. + c_1d_cutile_direct = vec_add( + a_1d_direct, b_1d_direct, use_gather=False, + num_cpu_threads=args.num_cpu_threads) + print( + f"""cuTile Output 1D shape: {c_1d_cutile_direct.shape}, + dtype: {c_1d_cutile_direct.dtype}""") + if args.correctness_check: + torch.testing.assert_close(c_1d_cutile_direct, a_1d_direct + b_1d_direct) + print("Correctness check passed") + else: + print("Correctness check disabled") + + # --- Test Case 2: 1D Vector Add (Gather/Scatter) --- + print("\n--- Test 2: 1D Vector Add (Gather/Scatter) ---") + # Use a size not perfectly divisible by typical TILE_SIZE to demonstrate + # the gather/scatter kernel's robust boundary handling. + VECTOR_SIZE_1D_GATHER = 1_000_001 + a_1d_gather = torch.randn(VECTOR_SIZE_1D_GATHER, dtype=torch.float32, device='cpu') + b_1d_gather = torch.randn(VECTOR_SIZE_1D_GATHER, dtype=torch.float32, device='cpu') + print(f"Input 1D (gather) shape: {a_1d_gather.shape}, dtype: {a_1d_gather.dtype}") + + # Call the vec_add wrapper function, requesting the gather/scatter kernel. + c_1d_cutile_gather = vec_add( + a_1d_gather, b_1d_gather, use_gather=True, + num_cpu_threads=args.num_cpu_threads) + print( + f"""cuTile Output 1D (gather) shape: {c_1d_cutile_gather.shape}, + dtype: {c_1d_cutile_gather.dtype}""") + if args.correctness_check: + torch.testing.assert_close(c_1d_cutile_gather, a_1d_gather + b_1d_gather) + print("Correctness check passed") + else: + print("Correctness check disabled") + + # --- Test Case 3: 2D Matrix Add (Direct Tiled) --- + print("\n--- Test 3: 2D Matrix Add (Direct Tiled) ---") + a_2d_direct = torch.randn(MATRIX_SHAPE_2D, dtype=torch.float32, device='cpu') + b_2d_direct = torch.randn(MATRIX_SHAPE_2D, dtype=torch.float32, device='cpu') + print(f"Input 2D shape: {a_2d_direct.shape}, dtype: {a_2d_direct.dtype}") + + # Call the vec_add wrapper function for 2D, requesting the direct tiled kernel. + c_2d_cutile_direct = vec_add( + a_2d_direct, b_2d_direct, use_gather=False, + num_cpu_threads=args.num_cpu_threads) + print( + f"""cuTile Output 2D shape: {c_2d_cutile_direct.shape}, + dtype: {c_2d_cutile_direct.dtype}""") + if args.correctness_check: + torch.testing.assert_close(c_2d_cutile_direct, a_2d_direct + b_2d_direct) + print("Correctness check passed") + else: + print("Correctness check disabled") + + # --- Test Case 4: 2D Matrix Add (Gather/Scatter) --- + print("\n--- Test 4: 2D Matrix Add (Gather/Scatter) ---") + # Use dimensions not perfectly divisible by typical tile sizes to demonstrate + # the gather/scatter kernel's robust boundary handling in 2D. + MATRIX_SHAPE_2D_GATHER = (2000, 1000) + a_2d_gather = torch.randn(MATRIX_SHAPE_2D_GATHER, dtype=torch.float32, device='cpu') + b_2d_gather = torch.randn(MATRIX_SHAPE_2D_GATHER, dtype=torch.float32, device='cpu') + print(f"Input 2D (gather) shape: {a_2d_gather.shape}, dtype: {a_2d_gather.dtype}") + + # Call the vec_add wrapper function for 2D, requesting the gather/scatter kernel. + c_2d_cutile_gather = vec_add( + a_2d_gather, b_2d_gather, use_gather=True, + num_cpu_threads=args.num_cpu_threads) + print( + f"""cuTile Output 2D (gather) shape: {c_2d_cutile_gather.shape}, + dtype: {c_2d_cutile_gather.dtype}""") + if args.correctness_check: + torch.testing.assert_close(c_2d_cutile_gather, a_2d_gather + b_2d_gather) + print("Correctness check passed") + else: + print("Correctness check disabled") + + print("\n--- cuTile Vector/Matrix Addition examples complete ---") diff --git a/samples/cpu/utils/benchmark.py b/samples/cpu/utils/benchmark.py new file mode 100644 index 00000000..9ddb5c7f --- /dev/null +++ b/samples/cpu/utils/benchmark.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import gc +import time +from math import ceil + +import cuda.tile as ct +from cuda.tile._backends import compile_for_launch, launch_compiled + + +_WARMUP_ITER_GUESS = 5 +_MIN_ROUND_TIME_NS = 200_000_000 +_ROUNDS = 50 +_WARMUP_ROUNDS = 1 +_MAX_ITERS = 1_000_000 + + +def _measure_block_ns(f, tuple_of_args, iterations: int) -> int: + gc_was_enabled = gc.isenabled() + if gc_was_enabled: + gc.disable() + try: + start = time.perf_counter_ns() + for _ in range(iterations): + f(*tuple_of_args) + return time.perf_counter_ns() - start + finally: + if gc_was_enabled: + gc.enable() + + +def _estimate_bench_iter(f, tuple_of_args): + f(*tuple_of_args) + elapsed_ns = _measure_block_ns( + f, tuple_of_args, _WARMUP_ITER_GUESS) + if elapsed_ns <= 0: + raise RuntimeError("CPU benchmark timer did not advance") + + iterations = ceil( + _MIN_ROUND_TIME_NS * _WARMUP_ITER_GUESS / elapsed_ns) + iterations = min(_MAX_ITERS, max(_ROUNDS, iterations)) + return _WARMUP_ROUNDS, iterations, _ROUNDS + + +def _time_ms(f, tuple_of_args, warmup: int, iters: int, rounds: int) -> float: + for _ in range(warmup): + f(*tuple_of_args) + + run_iters = max(iters, rounds) + elapsed_ns = sum( + _measure_block_ns(f, tuple_of_args, run_iters) + for _ in range(rounds) + ) + return elapsed_ns / (rounds * run_iters * 1_000_000) + + +def _report_benchmark(f, tuple_of_args) -> dict[str, float]: + warmup_rounds, iterations, rounds = _estimate_bench_iter( + f, tuple_of_args) + mean_time_ms = _time_ms( + f, tuple_of_args, warmup_rounds, iterations, rounds) + return {"mean_time_ms": mean_time_ms} + + +def report_benchmark(f, tuple_of_args, *, kernel=None, grid=None, + kernel_args=None, options=None) -> dict[str, float]: + """Benchmark a callable, optionally using one resolved CPU kernel binary. + + When ``kernel`` is provided, ``kernel_args`` are used to resolve the + backend binary once and each measured iteration launches that binary + directly. The original callable and ``tuple_of_args`` remain the regular + benchmark interface for callers that do not need this optimization. + """ + if kernel is None: + return _report_benchmark(f, tuple_of_args) + if grid is None or kernel_args is None: + raise TypeError( + "grid and kernel_args are required when kernel is provided") + + kernel_args = tuple(kernel_args) + with ct.backend_options(dict(options or {})): + compiled = compile_for_launch(kernel, kernel_args) + + def direct_call(): + return launch_compiled(None, grid, compiled, kernel_args) + + return _report_benchmark(direct_call, ()) diff --git a/setup.py b/setup.py index edff2e53..4df72bc4 100644 --- a/setup.py +++ b/setup.py @@ -6,12 +6,43 @@ from setuptools.command.build_ext import build_ext from setuptools.extension import Extension import os +import shlex import sys +import tempfile -project_root = os.path.dirname(os.path.abspath(__file__)) +project_root = os.path.dirname(os.path.realpath(__file__)) is_windows = sys.platform == "win32" +# Every build variant builds the _cext extension. The non-CUDA variants (see +# backends/README.md) also build their backend from backends/ and stage it, with +# its runtime files, into cuda/tile/_backends with `cmake --install`. +_backend_extensions = { + "cpu": "cuda.tile._backends.cpu._cpu", +} +build_variant = os.getenv("CUTILE_BUILD_VARIANT", "cuda") +if build_variant != "cuda" and build_variant not in _backend_extensions: + raise RuntimeError("CUTILE_BUILD_VARIANT must be one of: " + f"{', '.join(['cuda', *_backend_extensions])}") + + +def _backend_cmake_args(): + """CMake arguments of the non-CUDA variants, including those from CMAKE_ARGS.""" + args = shlex.split(os.getenv("CMAKE_ARGS", "")) + if "CUDAToolkit_ROOT" in os.environ: + return args + # Without a CUDA Toolkit, take the headers _cext needs from the + # cuda-toolkit[cudart,nvvm] wheels. + try: + import nvidia + except ImportError: + return args + for path in nvidia.__path__: + root = os.path.join(path, "cu13") + if os.path.isfile(os.path.join(root, "include", "cuda.h")): + return [f"-DCUDAToolkit_ROOT={root}", *args] + return args + class BuildExtWithCmake(build_ext): user_options = build_ext.user_options + [ @@ -52,7 +83,10 @@ def _cmake(self, build_dir: str, build_type: str, dlpack_path: str, xla_path: st f"-DXLA_PATH={xla_path}", f"-DCMAKE_BUILD_TYPE={build_type}", f"-DPython_EXECUTABLE={sys.executable}", - "-DCMAKE_POLICY_VERSION_MINIMUM=3.5"] + "-DCMAKE_POLICY_VERSION_MINIMUM=3.5", + f"-DCUTILE_BUILD_VARIANT={build_variant}"] + if build_variant != "cuda": + cmake_cmd.extend(_backend_cmake_args()) if self.disable_internal: cmake_cmd.append("-DDISABLE_INTERNAL=1") if self.minimal_internal: @@ -82,7 +116,14 @@ def run(self): self._cmake(build_dir, build_type, dlpack_path, xla_path) self._make(build_dir, build_type, parallel) + if build_variant != "cuda": + prefix = os.path.join(project_root, "src") if self.editable_mode else self.build_lib + self.spawn(["cmake", "--install", build_dir, "--prefix", os.path.abspath(prefix), + "--component", "cuTileBackend"]) + for ext in self.extensions: + if ext.name in _backend_extensions.values(): + continue # installed by `cmake --install` above src_dir = _get_csrc_dir(ext.name) ext_name = _get_build_lib_filename(ext.name) if is_windows: @@ -90,9 +131,22 @@ def run(self): else: ext_build_path = os.path.join(build_dir, src_dir, ext_name) ext_path = self.get_ext_fullpath(ext.name) - # Create a symlink to the build directory if in editable mode, otherwise copy - link = "sym" if self.editable_mode else None - file_util.copy_file(ext_build_path, ext_path, update=1, link=link) + if self.editable_mode: + _link_editable_extension(ext_build_path, ext_path) + else: + file_util.copy_file(ext_build_path, ext_path, update=1) + + +def _link_editable_extension(src: str, dst: str): + src = os.path.abspath(src) + if not os.path.isfile(src): + raise FileNotFoundError(src) + if os.path.islink(dst) and os.readlink(dst) == src: + return + with tempfile.TemporaryDirectory(dir=os.path.dirname(dst)) as temp_dir: + link = os.path.join(temp_dir, os.path.basename(dst)) + os.symlink(src, link) + os.replace(link, dst) def _get_csrc_dir(ext_name: str): @@ -110,8 +164,12 @@ def _get_build_lib_filename(ext_name: str): setup( + options=({"build": {"build_base": "build/python-" + build_variant}} + if build_variant != "cuda" else {}), ext_modules=[ Extension("cuda.tile._cext", []), + *([Extension(_backend_extensions[build_variant], [])] + if build_variant != "cuda" else []), ], cmdclass=dict( build_ext=BuildExtWithCmake, diff --git a/src/cuda/tile/__init__.py b/src/cuda/tile/__init__.py index 87d75d85..11da892c 100644 --- a/src/cuda/tile/__init__.py +++ b/src/cuda/tile/__init__.py @@ -4,8 +4,6 @@ from cuda.tile._version import __version__ # noqa -from cuda.tile._cext import launch - from cuda.tile._by_target import ByTarget from cuda.tile._memory_model import ( @@ -179,6 +177,9 @@ from cuda.tile._context import compiler_timeout from cuda.tile import tune +from cuda.tile.tune import autotune + +from cuda.tile._backends import launch, set_backend, backend_options from cuda.tile._execution import ( function, @@ -189,6 +190,8 @@ __all__ = [ "launch", + "set_backend", + "backend_options", "ByTarget", @@ -354,6 +357,7 @@ "compiler_timeout", "tune", + "autotune", "function", "kernel", diff --git a/src/cuda/tile/_backends/.gitignore b/src/cuda/tile/_backends/.gitignore new file mode 100644 index 00000000..997d5f00 --- /dev/null +++ b/src/cuda/tile/_backends/.gitignore @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# Runtime files staged by `cmake --install --prefix src`; extension +# modules are covered by the repository-wide `*.so` rule. +/_lib/ +/_licenses/ diff --git a/src/cuda/tile/_backends/__init__.py b/src/cuda/tile/_backends/__init__.py new file mode 100644 index 00000000..1e00a269 --- /dev/null +++ b/src/cuda/tile/_backends/__init__.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 +"""Pluggable backends for running cuTile kernels on non-CUDA targets. + +cuTile compiles a kernel's Python source to TileIR bytecode. A backend turns +that bytecode into a binary for its target and launches it. Select one with +:func:`cuda.tile.set_backend`, e.g. ``ct.set_backend("cpu")``; afterwards +:func:`cuda.tile.launch` compiles and runs kernels on it, using the options of +the enclosing :func:`cuda.tile.backend_options` block. Signature derivation, +in-memory and on-disk caching are shared by all backends. + +A backend is a module or object implementing the :class:`Backend` protocol. +Selection is process-wide and independent of arguments. Call +``cuda.tile.set_backend(None)`` to launch on CUDA. +""" + +from ._custom import ( + Backend, + CompiledKernel, + backend_options, + benchmark, + benchmark_callable, + compile_for_launch, + current_options, + get_backend, + launch, + launch_compiled, + set_backend, +) +from ._signature import ArrayArgument + +__all__ = ( + "ArrayArgument", + "Backend", + "CompiledKernel", + "backend_options", + "benchmark", + "benchmark_callable", + "compile_for_launch", + "current_options", + "get_backend", + "launch", + "launch_compiled", + "set_backend", +) diff --git a/src/cuda/tile/_backends/_custom.py b/src/cuda/tile/_backends/_custom.py new file mode 100644 index 00000000..9e650c28 --- /dev/null +++ b/src/cuda/tile/_backends/_custom.py @@ -0,0 +1,517 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 +"""Backend registry and launch dispatch for non-CUDA backends.""" + +from __future__ import annotations + +from collections import OrderedDict +import contextlib +from dataclasses import dataclass +import gc +import importlib +import logging +import sys +import threading +import time +from types import MappingProxyType +from typing import Any, Hashable, Mapping, Protocol + +from cuda.tile._cext import _synchronize_context, default_tile_context +from cuda.tile._cext import launch as _cext_launch +from cuda.tile.compilation import ConstantConstraint, KernelSignature, mangle_kernel_name +from cutile_cache._cache import cache_key, cache_lookup, cache_store, evict_lru + +from ._signature import ( + ArrayArgument, array_metadata, build_signature, constant_key, constant_parameters, +) + +logger = logging.getLogger(__name__) + + +class Backend(Protocol): + """A backend module or object, as accepted by :func:`set_backend`.""" + + #: Target name, e.g. ``"cpu"``; part of the disk cache key. + name: str + #: TileIR bytecode version that ``compile`` reads, e.g. ``"13.3"``. + bytecode_version: str + + def normalize_options(self, options: Mapping[str, Any]) -> Hashable: + """Validate the options of a :func:`backend_options` block, rejecting + unknown keys, into a value with a stable ``repr`` that covers everything + affecting ``compile``.""" + + def compiler_identity(self, options: Hashable) -> str | None: + """Identify everything besides the bytecode that affects ``compile``; + ``None`` disables the disk cache.""" + + def compile(self, bytecode: bytes, signature: KernelSignature, + options: Hashable) -> bytes: + """Compile TileIR bytecode into a binary.""" + + def load(self, binary: bytes, signature: KernelSignature, options: Hashable) -> Any: + """Prepare a binary for launching; called once per binary.""" + + def launch(self, loaded: Any, stream: Any, grid: Any, args: tuple) -> None: + """Run ``load``'s result synchronously. ``args`` holds the kernel arguments + without compile-time constants, with each array as an :class:`ArrayArgument`. + + A backend may also define ``launch_flat(loaded, stream, grid, values)``, + which receives each array as its pointer, shape and strides, inline. It is + used only for argument layouts that ``launch`` has accepted before. + + A backend may also define ``make_launch_cache(constants, limit)`` to + provide a native cache for repeated launches. Its ``launch`` returns + ``True`` on execution, ``None`` for unsupported arguments, or an opaque + layout token on a miss. ``insert(args, loaded, token)`` must reject + layouts that changed since that lookup.""" + + def synchronize(self, stream: Any) -> None: + """Wait for work queued on ``stream``.""" + + +_BACKEND_ATTRIBUTES = (*Backend.__annotations__, + *(name for name, value in vars(Backend).items() + if callable(value) and not name.startswith("_"))) + +# Per kernel, at most this many argument layouts and binaries stay cached. +_CACHE_LIMIT = 256 + +_NO_OPTIONS = MappingProxyType({}) + +_backend = None + + +class _ThreadState(threading.local): + options = None + normalized = None + + +_tls = _ThreadState() + + +def _import_backend_module(module: str): + # Support shorthand names (e.g. "cpu") for built-in backends. + candidates = [module] + if "." not in module: + candidates.insert(0, f"cuda.tile._backends.{module}") + + for name in candidates: + try: + return importlib.import_module(name) + except ModuleNotFoundError as exc: + # Only ignore "module not found" for the candidate itself. + if exc.name != name: + raise + raise ModuleNotFoundError( + f"Could not import backend module '{module}'. " + f"Tried: {', '.join(candidates)}" + ) + + +def set_backend(backend) -> None: + """Select the backend that runs :func:`cuda.tile.launch`. + + The selection is process-wide and is not inferred from kernel arguments. + Callers must coordinate backend changes across threads. + :func:`backend_options` blocks apply to the current thread. + + Args: + backend: A :class:`Backend` module (or object) or its name. Short names + such as ``"cpu"`` select the built-in backends. ``None`` + restores the default CUDA path. + """ + global _backend + if isinstance(backend, str): + backend = _import_backend_module(backend) + if backend is not None: + missing = [name for name in _BACKEND_ATTRIBUTES + if not hasattr(backend, name)] + if missing: + raise TypeError( + f"{backend!r} is not a cuTile backend; it lacks " + f"{', '.join(missing)}") + hash(backend) # identifies cached binaries + _backend = backend + + +def get_backend(): + """Return the active backend, or ``None`` for CUDA.""" + return _backend + + +@contextlib.contextmanager +def backend_options(options): + """Scope options of the active backend for the compiles and launches in the block. + + The backend's ``normalize_options`` validates them. Options that affect + code generation select the compiled binary; the others, such as the CPU + backend's ``num_cpu_threads``, only apply when launching. + """ + previous = getattr(_tls, "options", None) + _tls.options = dict(options or {}) + try: + yield + finally: + _tls.options = previous + + +def current_options(): + """Options of the innermost enclosing :func:`backend_options` block.""" + return getattr(_tls, "options", None) or _NO_OPTIONS + + +def _normalized_options(backend): + options = current_options() + memo = getattr(_tls, "normalized", None) + if memo is not None and memo[0] is backend and memo[1] is options: + return memo[2] + normalized = backend.normalize_options(options) + _tls.normalized = (backend, options, normalized) + return normalized + + +@dataclass(frozen=True, eq=False) +class CompiledKernel: + """A backend binary with everything needed to launch it.""" + + backend: Any + binary: bytes + signature: KernelSignature + options: Hashable + loaded: Any + constants: tuple[bool, ...] + + @property + def symbol(self) -> str: + return self.signature.symbol + + +class _KernelCache: + __slots__ = ("constants", "launches", "flat_launches", "native_launch", + "native_launches", "binaries", "lock", + "compile_lock") + + def __init__(self, constants: tuple[bool, ...]): + # Whether each kernel parameter is a compile-time constant. + self.constants = constants + # (backend, options, builder, argument layouts) -> CompiledKernel. + self.launches = OrderedDict() + # Layouts that `launch` accepted, keyed by _flat_arguments -> CompiledKernel. + self.flat_launches = OrderedDict() + self.native_launch = None + self.native_launches = OrderedDict() + # (backend, options, mangled name, symbol) -> CompiledKernel; only + # accessed under compile_lock. + self.binaries = OrderedDict() + # Guards `launches`; held only briefly so that cache hits never wait + # for a compile. + self.lock = threading.Lock() + # Serializes the compiles of this kernel. + self.compile_lock = threading.Lock() + + +def _kernel_cache(kernel) -> _KernelCache: + try: + return kernel._cutile_backend_cache + except AttributeError: + return kernel.__dict__.setdefault( + "_cutile_backend_cache", _KernelCache(constant_parameters(kernel))) + + +def _lru_get(cache: OrderedDict, key): + value = cache.get(key) + if value is not None: + cache.move_to_end(key) + return value + + +def _lru_put(cache: OrderedDict, key, value) -> None: + cache[key] = value + if len(cache) > _CACHE_LIMIT: + cache.popitem(last=False) + + +def _describe(constants, args) -> tuple: + """Return each argument's :class:`ArrayArgument`, or ``None`` if it is none.""" + if len(args) != len(constants): + raise TypeError( + f"kernel expects {len(constants)} arguments, got {len(args)}") + arrays = [] + for value, constant in zip(args, constants): + array = None + if not constant: + try: + array = array_metadata(value) + except TypeError: + pass + arrays.append(array) + return tuple(arrays) + + +def _argument_key(value, constant: bool, array: ArrayArgument | None): + """Everything about ``value`` that a signature builder may depend on.""" + if constant: + return constant_key(value) + if array is None: + return type(value) + pointer = array.pointer + return array.dtype, array.shape, array.strides, pointer & -pointer + + +_SCALAR_TYPES = (int, float, bool) + + +def _flat_arguments(backend, constants, args): + """Return a layout key and ``launch_flat`` values for ``args``, or ``None`` if + an argument is neither a builtin scalar nor an array with ``data_ptr()``. + + The key covers everything ``_argument_key`` and ``check_argument`` depend on. + """ + if len(args) != len(constants): + return None + key = [backend, _normalized_options(backend)] + values = [] + tensor_type = getattr(sys.modules.get("torch"), "Tensor", None) + for value, constant in zip(args, constants): + if constant: + key.append(constant_key(value)) + elif type(value) in _SCALAR_TYPES: + key.append(type(value)) + values.append(value) + else: + data_ptr = getattr(value, "data_ptr", None) + if not callable(data_ptr): + return None + pointer = int(data_ptr()) + if type(value) is tensor_type: + shape = value.shape + strides = value.stride() + else: + shape = tuple(map(int, value.shape)) + strides = tuple(map(int, value.stride())) + key.append((value.dtype, getattr(value, "device", "cpu"), shape, strides, + pointer & -pointer)) + values.append(pointer) + values += shape + values += strides + return tuple(key), values + + +def _runtime_arguments(args, constants, arrays) -> tuple: + """The arguments a backend's ``launch`` receives; see :class:`Backend`.""" + return tuple(value if array is None else array + for value, constant, array in zip(args, constants, arrays) + if not constant) + + +def _compile_for_launch(backend, kernel, cache, args, arrays, signature_builder): + options = _normalized_options(backend) + key = (backend, options, signature_builder, + tuple(map(_argument_key, args, cache.constants, arrays)), + # A custom builder may assume divisibility by a non-power of two. + tuple(array.pointer if array is not None else None for array in arrays) + if signature_builder is not build_signature else None) + with cache.lock: + compiled = _lru_get(cache.launches, key) + if compiled is not None: + return compiled + with cache.compile_lock: + with cache.lock: + compiled = _lru_get(cache.launches, key) + if compiled is None: + compiled = _compile(kernel, cache, backend, + signature_builder(kernel, args), options) + with cache.lock: + _lru_put(cache.launches, key, compiled) + return compiled + + +def _active_backend(): + backend = _backend + if backend is None: + raise RuntimeError("no backend is active; see cuda.tile.set_backend()") + return backend + + +def compile_for_launch(kernel, args, *, signature_builder=build_signature): + """Return the active backend's :class:`CompiledKernel` for ``args``. + + ``signature_builder(kernel, args)`` may add assumptions to the signature, + as long as they only depend on constants, scalar types and array + dtypes, shapes, strides and alignment. + """ + backend = _active_backend() + args = tuple(args) + cache = _kernel_cache(kernel) + arrays = _describe(cache.constants, args) + return _compile_for_launch( + backend, kernel, cache, args, arrays, signature_builder) + + +def _compile(kernel, cache, backend, signature, options): + name = kernel._annotated_function.pyfunc.__name__ + key = (backend, options, mangle_kernel_name(name, signature), signature.symbol) + compiled = _lru_get(cache.binaries, key) + if compiled is None: + binary, signature = _compile_binary(kernel, backend, signature, options) + loaded = backend.load(binary, signature, options) + constants = tuple(isinstance(parameter, ConstantConstraint) + for parameter in signature.parameters) + compiled = CompiledKernel( + backend, binary, signature, options, loaded, constants) + _lru_put(cache.binaries, key, compiled) + return compiled + + +def compile_bytecode(kernel, backend, signature) -> tuple[bytes, KernelSignature]: + """Compile ``kernel`` to architecture-independent TileIR bytecode for ``backend``.""" + from cuda.tile._compile import compile_tile, parse_bytecode_version + + result = compile_tile( + kernel._annotated_function, (signature,), None, + kernel._compiler_options, default_tile_context, + bytecode_version=parse_bytecode_version(backend.bytecode_version), + return_bytecode=True, return_cubin=False) + [signature] = result.kernel_signatures + return bytes(result.bytecode), signature + + +def _compile_binary(kernel, backend, signature, options): + bytecode, signature = compile_bytecode(kernel, backend, signature) + + context = default_tile_context + cache_dir = context.config.cache_dir + disk_key = None + if cache_dir is not None: + identity = backend.compiler_identity(options) + if identity is None: + logger.warning("disk cache disabled: backend compiler identity is unknown") + else: + # Backends ignore the kernel's opt_level, so it is not part of the key. + disk_key = cache_key(f"cutile-backend-v3:{identity}", backend.name, 0, bytecode) + binary = cache_lookup(cache_dir, disk_key) + if binary is not None: + return binary, signature + + binary = backend.compile(bytecode, signature, options) + if not isinstance(binary, bytes): + raise TypeError("backend compile() must return bytes") + if disk_key is not None: + cache_store(cache_dir, disk_key, binary) + evict_lru(cache_dir, context.config.cache_size_limit) + return binary, signature + + +def launch(stream, grid, kernel, kernel_args, /, **options): + """Launch a cuTile kernel. + + Runs on the backend selected with :func:`cuda.tile.set_backend`, or on CUDA + if no backend is selected. Arguments never override the selection. + + Args: + stream: The stream to execute the |kernel| on. + grid: Tuple of up to 3 grid dimensions to execute the |kernel| over. + kernel: The |kernel| to execute. + kernel_args: Positional arguments to pass to the kernel. + **options: CUDA launch options, such as ``programmatic_dependent_launch``. + """ + backend = _backend + if backend is None: + # Calls without options avoid the slower keyword-argument call. + if options: + return _cext_launch(stream, grid, kernel, kernel_args, **options) + return _cext_launch(stream, grid, kernel, kernel_args) + if options: + raise TypeError(f"launch options {', '.join(map(repr, options))} are only " + f"supported on CUDA") + args = tuple(kernel_args) + cache = _kernel_cache(kernel) + native = None + factory = getattr(backend, "make_launch_cache", None) + if factory is not None: + native_key = (backend, _normalized_options(backend)) + memo = cache.native_launch + if memo is not None and memo[0] == native_key: + native = memo[1] + else: + with cache.lock: + native = _lru_get(cache.native_launches, native_key) + if native is None: + native = factory(cache.constants, _CACHE_LIMIT) + _lru_put(cache.native_launches, native_key, native) + if len(cache.native_launches) > 4: + cache.native_launches.popitem(last=False) + cache.native_launch = (native_key, native) + native_result = native.launch(args, stream, grid, current_options()) + if native_result is True: + return + if native_result is None: + native = None + launch_flat = getattr(backend, "launch_flat", None) + flat = launch_flat and _flat_arguments(backend, cache.constants, args) + if flat: + key, values = flat + try: + with cache.lock: + compiled = _lru_get(cache.flat_launches, key) + except TypeError: # an unhashable argument property + flat = None + else: + if compiled is not None: + launch_flat(compiled.loaded, stream, grid, values) + if native is not None: + native.insert(args, compiled.loaded, native_result) + return + arrays = _describe(cache.constants, args) + compiled = _compile_for_launch( + backend, kernel, cache, args, arrays, build_signature) + backend.launch(compiled.loaded, stream, grid, + _runtime_arguments(args, cache.constants, arrays)) + if native is not None: + native.insert(args, compiled.loaded, native_result) + if flat: + with cache.lock: + _lru_put(cache.flat_launches, key, compiled) + + +def launch_compiled(stream, grid, compiled: CompiledKernel, kernel_args, /): + """Launch ``compiled`` without dispatch lookups. + + ``kernel_args`` must match the arguments ``compiled`` was built for. + """ + args = tuple(kernel_args) + arrays = _describe(compiled.constants, args) + compiled.backend.launch(compiled.loaded, stream, grid, + _runtime_arguments(args, compiled.constants, arrays)) + + +def _synchronize(stream): + if _backend is None: + _synchronize_context() + else: + _backend.synchronize(stream) + + +def benchmark_callable(stream, fn, args=()) -> float: + """Time one synchronous call of ``fn(*args)`` in microseconds.""" + _synchronize(stream) + gc_was_enabled = gc.isenabled() + gc.disable() + try: + start = time.perf_counter_ns() + fn(*args) + _synchronize(stream) + elapsed_ns = time.perf_counter_ns() - start + finally: + if gc_was_enabled: + gc.enable() + return elapsed_ns / 1_000.0 + + +def benchmark(stream, grid, kernel, args) -> float: + """Time one launch on the active backend in microseconds.""" + compiled = compile_for_launch(kernel, args) + return benchmark_callable( + stream, launch_compiled, (stream, grid, compiled, args)) diff --git a/src/cuda/tile/_backends/_native.py b/src/cuda/tile/_backends/_native.py new file mode 100644 index 00000000..25855443 --- /dev/null +++ b/src/cuda/tile/_backends/_native.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 +"""The native compiler shared by the built-in backends. + +The private nanobind extension links the pinned LLVM/MLIR and TileIRToMLIR +libraries. In-process compilation releases the GIL and uses a fresh MLIR context +per call. An explicit compiler timeout selects a Python worker running the same +extension so a timeout can terminate compilation safely. Without a timeout, +native compiler crashes affect the calling process. + +The worker preserves known ValueError and RuntimeError failures without an +additional worker traceback. Installed kernels do not invoke LLVM tools or +require MLIR Python bindings. LLVM/MLIR assertions are disabled by default; +CMAKE_ARGS=-DCUTILE_LLVM_ENABLE_ASSERTIONS=ON enables them for debugging. + +Build configuration +------------------- +The CMake build fetches pinned TileIRToMLIR and CUDA Tile archives, verified +by SHA-256, like the XLA headers. TileIRToMLIR's llvm-revision.txt +defines the required LLVM commit. USE_LLVM_REVISION may select a newer +descendant, but older or divergent commits are rejected. CMAKE_ARGS accepts +CUTILE_LLVM_SOURCE_DIR and CUTILE_LLVM_BINARY_DIR for custom build locations. +CUDAToolkit_ROOT selects local CUDA headers instead of the build dependencies. + +To reuse a prebuilt compiler, pass +CMAKE_ARGS="-DCUTILE_TOOLCHAIN_DIR=/path/to/prefix" to the backend's uv sync +command. The prefix must match the required LLVM revision, active Python +version, and build variant. It must contain lib/_compiler and +share/cutile/llvm/LICENSE.TXT; the cuTileCompiler CMake install component +stages both. Reusing the prefix skips the LLVM checkout and build. Package +configuration still fetches the dependency sources for their license files. +""" + +import hashlib +import importlib +import json +import os +import subprocess +import sys + +from cuda.tile._cext import default_tile_context +from cuda.tile._exception import TileCompilerTimeoutError + +# TileIR bytecode version read by the native compiler's TileIRToMLIR. +BYTECODE_VERSION = "13.4" + + +def file_fingerprint(path: str) -> str: + """Identify a file by canonical path and stat data, so any rebuild changes it.""" + resolved = os.path.realpath(path) + stat = os.stat(resolved) + return (f"{resolved}:{stat.st_dev}:{stat.st_ino}:{stat.st_size}:" + f"{stat.st_mtime_ns}") + + +def identity_digest(identity: dict) -> str: + """Digest of a JSON-serializable compiler identity.""" + return hashlib.sha256( + json.dumps(identity, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + + +def compiler(backend: str): + module = importlib.import_module("cuda.tile._backends._compiler") + if module.backend != backend: + raise RuntimeError( + f"The installed native compiler targets {module.backend}, not {backend}; " + f"rebuild with CUTILE_BUILD_VARIANT={backend}.") + return module + + +# Loads only the compiler extension, not the cuda.tile package. +_WORKER = """\ +import importlib.machinery +import importlib.util +import json +import sys +path, method, arguments = json.loads(sys.argv[1]) +loader = importlib.machinery.ExtensionFileLoader("_compiler", path) +_compiler = importlib.util.module_from_spec( + importlib.util.spec_from_loader("_compiler", loader)) +loader.exec_module(_compiler) +try: + result = getattr(_compiler, method)(sys.stdin.buffer.read(), *arguments) +except (ValueError, RuntimeError) as error: + sys.stderr.write(str(error)) + sys.exit(2 if isinstance(error, ValueError) else 3) +sys.stdout.buffer.write(result) +""" + + +def invoke(backend: str, method: str, bytecode: bytes, *arguments) -> bytes: + module = compiler(backend) + timeout = default_tile_context.config.compiler_timeout_sec + if timeout is None: + return getattr(module, method)(bytecode, *arguments) + # An explicit hard timeout needs a killable worker, not a C++ thread. + try: + process = subprocess.run( + [sys.executable, "-c", _WORKER, + json.dumps((module.__file__, method, arguments))], + input=bytecode, capture_output=True, timeout=timeout, check=False) + except subprocess.TimeoutExpired: + raise TileCompilerTimeoutError( + f"{backend} native compiler exceeded {timeout}s.", + method, None) from None + if process.returncode in (2, 3): + exception = ValueError if process.returncode == 2 else RuntimeError + raise exception(process.stderr.decode(errors="replace")) + if process.returncode: + raise RuntimeError( + f"{backend} native compiler failed (exit code {process.returncode}).\n" + + process.stderr.decode(errors="replace")) + if process.stderr: + sys.stderr.write(process.stderr.decode(errors="replace")) + sys.stderr.flush() + return process.stdout diff --git a/src/cuda/tile/_backends/_signature.py b/src/cuda/tile/_backends/_signature.py new file mode 100644 index 00000000..e8005ca3 --- /dev/null +++ b/src/cuda/tile/_backends/_signature.py @@ -0,0 +1,264 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import struct +from typing import Any, NamedTuple + +import cuda.tile as ct +from cuda.tile._annotated_function import LeafAnnotationNode +from cuda.tile._ir.typing_support import to_dtype +from cuda.tile.compilation import ( + ArrayConstraint, + CallingConvention, + ConstantConstraint, + KernelSignature, + ScalarConstraint, +) + + +# The native launcher (cext/tile_kernel.cpp) derives layout assumptions, +# including static size-one shapes, for at most this many dimensions. +_MAX_SPECIALIZED_NDIM = 5 + +_CUTILE_PYTHON_V2 = CallingConvention.cutile_python_v2() + +# Struct format of each runtime scalar dtype, as passed to a backend kernel. +_SCALAR_FORMATS = { + ct.bool_: "?", ct.int8: "b", ct.uint8: "B", ct.int16: "h", ct.uint16: "H", + ct.int32: "i", ct.uint32: "I", ct.int64: "q", ct.uint64: "Q", + ct.float32: "f", ct.float64: "d", +} + + +def constant_key(value): + """Preserve the type and float bits used by ConstantConstraint and mangling.""" + return type(value), struct.pack("=d", value) if isinstance(value, float) else value + + +class ArrayArgument(NamedTuple): + """An array kernel argument as passed to a backend's ``launch``.""" + + pointer: int + shape: tuple[int, ...] + strides: tuple[int, ...] # in elements + dtype: Any + device: str # device type, e.g. "cpu" + device_index: int + + +def array_metadata(value: Any) -> ArrayArgument: + """Describe an array argument; raise ``TypeError`` for anything else.""" + data_ptr = getattr(value, "data_ptr", None) + if callable(data_ptr): + device = getattr(value, "device", "cpu") + return ArrayArgument(int(data_ptr()), tuple(map(int, value.shape)), + tuple(map(int, value.stride())), value.dtype, + getattr(device, "type", device), + getattr(device, "index", None) or 0) + if not hasattr(value, "__array_interface__"): + raise TypeError( + f"expected an array with data_ptr() or __array_interface__, got " + f"{type(value).__name__}") + import numpy as np + + value = np.asarray(value) + itemsize = value.dtype.itemsize + if any(stride % itemsize for stride in value.strides): + raise ValueError("array strides must be multiples of the element size") + strides = tuple(stride // itemsize for stride in value.strides) + return ArrayArgument(value.ctypes.data, value.shape, strides, value.dtype, "cpu", 0) + + +def scalar_format(backend: str, dtype) -> str: + """Struct format of a runtime scalar of ``dtype``.""" + try: + return _SCALAR_FORMATS[dtype] + except KeyError: + raise TypeError(f"{backend} backend: unsupported scalar dtype {dtype}") from None + + +def runtime_parameters(backend: str, signature: KernelSignature) -> tuple: + """The signature's parameters that receive runtime arguments.""" + parameters = tuple(parameter for parameter in signature.parameters + if not isinstance(parameter, ConstantConstraint)) + for parameter in parameters: + if isinstance(parameter, ScalarConstraint): + scalar_format(backend, parameter.dtype) + elif not isinstance(parameter, ArrayConstraint): + raise TypeError(f"{backend} backend: unsupported parameter constraint " + f"{type(parameter).__name__}") + return parameters + + +def check_argument(backend: str, device: str, parameter, argument) -> None: + """Check that a runtime argument matches its signature parameter.""" + if isinstance(parameter, ArrayConstraint): + if (not isinstance(argument, ArrayArgument) + or len(argument.shape) != parameter.ndim): + raise TypeError(f"{backend} backend: expected a {parameter.ndim}-D array, " + f"got {type(argument).__name__}") + if argument.device != device: + raise ValueError(f"{backend} backend: array is on device " + f"'{argument.device}', expected '{device}'") + elif isinstance(argument, ArrayArgument): + raise TypeError(f"{backend} backend: expected a scalar, got an array") + + +def _elements_disjoint(shape, strides) -> bool: + # Same rule as the native launcher: sorted by stride, every stride must be + # positive and step over the whole extent of the previous dimension. + dims = sorted(zip(strides, shape)) + if not dims: + return True + return dims[0][0] > 0 and all( + stride > 0 and stride >= prev_stride * prev_size + for (prev_stride, prev_size), (stride, _) in zip(dims, dims[1:])) + + +def _normalized_dims(dims, ndim: int, name: str) -> tuple[int, ...]: + normalized = [] + for dim in dims: + if not -ndim <= dim < ndim: + raise ValueError(f"{name} axis {dim} is out of range for a {ndim}-D array") + dim %= ndim + if dim in normalized: + raise ValueError(f"{name} contains duplicate axis {dim}") + normalized.append(dim) + return tuple(normalized) + + +def _array_constraint(array: ArrayArgument, annotation, *, base_addr_divisible_by: int, + infer_shape_one: bool): + try: + dtype = to_dtype(array.dtype) + except KeyError: + raise TypeError(f"unsupported array dtype {array.dtype}") from None + + shape, strides = array.shape, array.strides + ndim = len(shape) + index_dtype = ct.int32 if annotation is None else annotation.index_dtype + static_shape_dims = () if annotation is None else _normalized_dims( + annotation.static_shape_dims, ndim, "static_shape_dims") + static_stride_dims = () if annotation is None else _normalized_dims( + annotation.static_stride_dims, ndim, "static_stride_dims") + if index_dtype is ct.int32: + i32_max = (1 << 31) - 1 + if any(size > i32_max or stride > i32_max + for size, stride in zip(shape, strides)): + raise TypeError( + "array shape or stride exceeds int32; annotate the parameter " + "with ct.int64 index dtype") + if any(stride < 0 for stride in strides): + raise ValueError("negative array strides are not supported") + + specialize = ndim <= _MAX_SPECIALIZED_NDIM + align_bits = 16 * 8 + stride_divisor = (align_bits // dtype.bitwidth + if align_bits % dtype.bitwidth == 0 else 1) + # As in the native launcher, annotated static strides and shapes replace the + # inference of unit strides and size-one shapes, respectively. + if static_stride_dims: + stride_constant = tuple(strides[i] if i in static_stride_dims else None + for i in range(ndim)) + else: + stride_constant = tuple( + 1 if specialize and stride == 1 else None for stride in strides) + return ArrayConstraint( + dtype=dtype, + ndim=ndim, + index_dtype=index_dtype, + stride_lower_bound_incl=0, + alias_groups=(), + may_alias_internally=not _elements_disjoint(shape, strides), + stride_constant=stride_constant, + shape_constant=tuple( + shape[i] if i in static_shape_dims + else 1 if (infer_shape_one and specialize and not static_shape_dims + and shape[i] == 1) + else None for i in range(ndim)), + stride_divisible_by=tuple( + stride_divisor if specialize and stride % stride_divisor == 0 else 1 + for stride in strides), + shape_divisible_by=tuple( + 16 if specialize and size % 16 == 0 else 1 for size in shape), + base_addr_divisible_by=(base_addr_divisible_by + if array.pointer % base_addr_divisible_by == 0 else 1), + ) + + +def _scalar_constraint(value: Any, annotation) -> ScalarConstraint: + if isinstance(value, bool): + return ScalarConstraint(ct.bool_) + if isinstance(value, int): + return ScalarConstraint(ct.int32 if annotation is None else annotation.dtype) + if isinstance(value, float): + return ScalarConstraint(ct.float32) + if isinstance(value, (list, tuple)): + raise TypeError("array list arguments are only supported on CUDA") + raise TypeError(f"unsupported kernel argument type: {type(value).__name__}") + + +def _leaf_annotations(kernel) -> tuple[LeafAnnotationNode, ...]: + annotations = kernel._annotated_function.parameter_annotations + for annotation in annotations: + if not isinstance(annotation, LeafAnnotationNode): + raise TypeError("tuple parameters are only supported on CUDA") + if annotation.list is not None: + raise TypeError("array list parameters are only supported on CUDA") + return annotations + + +def constant_parameters(kernel) -> tuple[bool, ...]: + """Whether each parameter of ``kernel`` is a ``ct.Constant`` scalar or array.""" + return tuple(isinstance(annotation, LeafAnnotationNode) and annotation.constant + for annotation in kernel._annotated_function.parameter_annotations) + + +def _requires_cutile_python_v2(parameter) -> bool: + return (isinstance(parameter, ArrayConstraint) + and any(size is not None for size in parameter.shape_constant)) + + +def build_signature(kernel, args, *, calling_convention=None, + base_addr_divisible_by=16, symbol=None) -> KernelSignature: + """The signature the CUDA launcher derives for ``args``, for any backend. + + Without ``calling_convention``, as for CUDA launches, size-one dimensions + become static shapes and the lowest sufficient calling convention is used. + Otherwise this matches ``KernelSignature.from_kernel_args``, which infers + static shapes only for ``cutile_python_v2`` and later. + """ + infer_shape_one = (calling_convention is None + or calling_convention.version >= _CUTILE_PYTHON_V2.version) + annotations = _leaf_annotations(kernel) + if len(args) != len(annotations): + raise TypeError( + f"kernel expects {len(annotations)} arguments, got {len(args)}") + + parameters = [] + for annotation, value in zip(annotations, args): + if annotation.constant: + parameters.append(ConstantConstraint(value)) + continue + try: + array = array_metadata(value) + except TypeError: + parameters.append(_scalar_constraint(value, annotation.scalar)) + else: + parameters.append(_array_constraint( + array, annotation.array, + base_addr_divisible_by=base_addr_divisible_by, + infer_shape_one=infer_shape_one, + )) + + convention = calling_convention or ( + _CUTILE_PYTHON_V2 if any(map(_requires_cutile_python_v2, parameters)) + else CallingConvention.cutile_python_v1()) + signature = KernelSignature(parameters, convention, symbol) + if symbol is None: + signature = signature.with_mangled_symbol( + kernel._annotated_function.pyfunc.__name__) + return signature diff --git a/src/cuda/tile/_backends/_util.py b/src/cuda/tile/_backends/_util.py new file mode 100644 index 00000000..69194926 --- /dev/null +++ b/src/cuda/tile/_backends/_util.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 +"""Helpers shared by the built-in backends.""" + + +def synchronize_stream(stream) -> None: + """Wait for a framework stream; ``None`` and ``0`` stand for no stream.""" + if stream is None or (isinstance(stream, int) and stream == 0): + return + synchronize = getattr(stream, "synchronize", None) + if synchronize is None: + raise TypeError("stream must be None, 0, or provide synchronize()") + synchronize() + + +def normalize_dims(value) -> tuple[int, int, int]: + """Normalize one to three launch dimensions to a three-tuple.""" + + values = (value,) if isinstance(value, int) else tuple(value) + if not 1 <= len(values) <= 3: + raise ValueError("launch dimensions must have length 1 to 3") + return values + (1,) * (3 - len(values)) + + +def check_options(backend: str, options, known) -> None: + """Reject option keys the backend does not know, e.g. misspelled ones.""" + unknown = sorted(set(options) - set(known)) + if unknown: + raise ValueError( + f"{backend} backend: unknown options {', '.join(map(repr, unknown))}; " + f"expected any of {', '.join(sorted(known))}") + + +def bool_option(backend: str, options, key: str, default: bool) -> bool: + """Read a boolean option, rejecting look-alikes such as ``"false"``.""" + + value = options.get(key, default) + if not isinstance(value, bool): + raise TypeError( + f"{backend} backend: option {key!r} must be a bool, got {value!r}") + return value diff --git a/src/cuda/tile/_backends/cpu/__init__.py b/src/cuda/tile/_backends/cpu/__init__.py new file mode 100644 index 00000000..f18836f0 --- /dev/null +++ b/src/cuda/tile/_backends/cpu/__init__.py @@ -0,0 +1,203 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 +"""CPU backend: native compilation and synchronous OpenMP execution. + +TileIRToMLIR and MLIR passes, LLVM optimization and object emission, and the +LLD ELF linker run through C++ APIs in the private native compiler. Shared +libraries use cuTile's binary cache. Installed kernels need no LLVM tools or +MLIR Python bindings. + +TileIRToMLIR emits a uniform ABI: a pointer to immutable eight-byte argument +slots followed by six i32 block-id/grid-dimension arguments. A single +nanobind/OpenMP launcher serves all kernels; no launcher is compiled per kernel. +Blocks use static OpenMP scheduling without a grid-sized coordinate array. +One worker and single-block grids bypass parallel dispatch. +The launcher uses the system GCC OpenMP runtime (libgomp.so.1), not a bundled copy. + +Before general lowering, a C++ pass removes leading unit vector +dimensions, including on transfer operations, on any host. It selects x86 +contraction-loop preparation on x86-64, and AMX lowering only when +AMX tile, BF16, and INT8 are available. Final LLVM conversion lowers AMX +operations and loop-carried tile types together with control flow. On Linux +x86-64, importing the launcher requests process-wide tile-data permission. +If permission is denied, pipeline selection uses the x86 path without AMX. +After vector preparation, math functions except sin and cos are replaced by +polynomial approximations so they stay vectorized; LLVM would otherwise +scalarize vector math intrinsics into per-element libm calls. + +The compiler and TileIRToMLIR are LLVM external projects in the persistent +build/llvm tree. Shared compiler invocation and timeout handling live in +cuda.tile._backends._native. See backends/cpu/README.md for build and usage. +""" + +from __future__ import annotations + +import dataclasses +import functools +import operator +import os +import platform +import struct +import tempfile +from typing import Any, Mapping, NamedTuple + +import cuda.tile as ct +from cuda.tile.compilation import ArrayConstraint + +from .._custom import current_options +from .._native import BYTECODE_VERSION +from .._native import compiler as native_compiler +from .._native import file_fingerprint, identity_digest, invoke +from .._signature import ArrayArgument, check_argument, runtime_parameters, scalar_format +from .._util import bool_option, check_options, normalize_dims +from .._util import synchronize_stream as synchronize + +name = "cpu" +bytecode_version = BYTECODE_VERSION + +# Options of a ``backend_options`` block; ``num_cpu_threads`` only applies at launch. +_OPTIONS = ("assume_in_bounds", "num_cpu_threads") + + +@dataclasses.dataclass(frozen=True) +class Options: + """Validated compile options.""" + + assume_in_bounds: bool + + +def normalize_options(options: Mapping[str, Any]) -> Options: + """Validate and normalize the options of a ``backend_options`` block.""" + check_options("CPU", options, _OPTIONS) + return Options( + assume_in_bounds=bool_option("CPU", options, "assume_in_bounds", False)) + + +@functools.lru_cache(maxsize=1) +def _host_cpu(): + """Include the host ISA in cache keys for native LLVM code generation.""" + with open("/proc/cpuinfo") as file: + fields = dict( + line.split(":", 1) for line in file.read().split("\n\n", 1)[0].splitlines() + if ":" in line) + return {key.strip(): value.strip() for key, value in fields.items() + if key.strip() in ("model name", "flags", "Features", "CPU architecture", + "CPU implementer", "CPU part", "CPU variant", "CPU revision")} + + +def _amx_enabled() -> bool: + """Whether the launcher enabled AMX tile, BF16 and INT8 for this process.""" + from ._cpu import amx_enabled + + return amx_enabled + + +def _transform_entry_point() -> str: + """The native compiler's built-in pipeline entry point for this host.""" + return native_compiler("cpu").cpu_transform_entry_point() + + +def compiler_identity(options: Options): + """Identify the native toolchain, host ISA, and effective compile options.""" + + try: + identity = { + "schema": 7, + "compiler": file_fingerprint(native_compiler("cpu").__file__), + "cutile_cpu_backend": file_fingerprint(__file__), + "entry_point": _transform_entry_point(), + "cpu": _host_cpu(), + "platform": (platform.system(), platform.machine()), + "libc": platform.libc_ver(), + "options": dataclasses.asdict(options), + } + except OSError: + return None + return identity_digest(identity) + + +def compile(bytecode: bytes, signature, options: Options) -> bytes: + """Compile TileIR into a shared object with the private native compiler.""" + del signature + with tempfile.TemporaryDirectory() as directory: + return invoke( + "cpu", "compile_cpu", bytecode, + options.assume_in_bounds, directory, + os.environ.get("CUTILE_CPU_DUMP_MLIR", ""), + os.environ.get("MLIR_ENABLE_DUMP") == "1") + + +def _argument_packer(formats) -> struct.Struct: + """Match TileIRToMLIR's eight-byte slots without scalar alignment gaps.""" + return struct.Struct("@" + "".join( + f"{fmt}{8 - struct.calcsize(fmt)}x" for fmt in formats)) + + +def _argument_formats(parameters) -> list[str]: + """Struct formats of the flat native ABI: a pointer, sizes and strides per array.""" + formats = [] + for parameter in parameters: + if isinstance(parameter, ArrayConstraint): + index = "q" if parameter.index_dtype is ct.int64 else "i" + formats += ["P", *([index] * parameter.ndim * 2)] + else: + formats.append(scalar_format("CPU", parameter.dtype)) + return formats + + +class _Loaded(NamedTuple): + parameters: tuple + kernel: Any + packer: struct.Struct + native_plan: tuple + + +def make_launch_cache(constants, limit): + from ._cpu import LaunchCache + + return LaunchCache(constants, limit) + + +def load(binary: bytes, signature, options: Options) -> _Loaded: + """Load the shared object through the generic nanobind/OpenMP launcher.""" + from ._cpu import Kernel + + del options + parameters = runtime_parameters("CPU", signature) + packer = _argument_packer(_argument_formats(parameters)) + native_plan = tuple( + (parameter.ndim, "q" if parameter.index_dtype is ct.int64 else "i") + if isinstance(parameter, ArrayConstraint) + else scalar_format("CPU", parameter.dtype) + for parameter in parameters) + # dlopen needs a file; TMPDIR selects where it goes. + with tempfile.NamedTemporaryFile(suffix=".so") as file: + file.write(binary) + file.flush() + kernel = Kernel(file.name, signature.symbol, packer.size) + return _Loaded(parameters, kernel, packer, native_plan) + + +def _flatten_arguments(parameters, args) -> list: + values = [] + for parameter, argument in zip(parameters, args, strict=True): + check_argument("CPU", "cpu", parameter, argument) + if isinstance(argument, ArrayArgument): + values += (argument.pointer, *argument.shape, *argument.strides) + else: + values.append(argument) + return values + + +def launch(loaded: _Loaded, stream, grid, args): + """Run a loaded kernel; the ``num_cpu_threads`` option sets the worker count.""" + launch_flat(loaded, stream, grid, _flatten_arguments(loaded.parameters, args)) + + +def launch_flat(loaded: _Loaded, stream, grid, values): + """Run a loaded kernel with arguments already flattened as by ``launch``.""" + synchronize(stream) + dimensions = tuple(operator.index(value) for value in normalize_dims(grid)) + threads = operator.index(current_options().get("num_cpu_threads", 0)) + loaded.kernel.launch(loaded.packer.pack(*values), dimensions, threads) diff --git a/src/cuda/tile/tune/__init__.py b/src/cuda/tile/tune/__init__.py index 9445d391..a42018f1 100644 --- a/src/cuda/tile/tune/__init__.py +++ b/src/cuda/tile/tune/__init__.py @@ -2,5 +2,19 @@ # # SPDX-License-Identifier: Apache-2.0 from cuda.tile.tune._tune import exhaustive_search, TuningResult, Measurement +from cuda.tile.tune._autotune import ( + Config, + AutotunedFunction, + AutotunedKernel, + autotune, +) -__all__ = ["exhaustive_search", "TuningResult", "Measurement"] +__all__ = [ + "exhaustive_search", + "TuningResult", + "Measurement", + "Config", + "AutotunedFunction", + "AutotunedKernel", + "autotune", +] diff --git a/src/cuda/tile/tune/_autotune.py b/src/cuda/tile/tune/_autotune.py new file mode 100644 index 00000000..fcc01c9f --- /dev/null +++ b/src/cuda/tile/tune/_autotune.py @@ -0,0 +1,583 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 +"""Triton-style ``@autotune`` for cuTile kernels and multi-kernel factories.""" +from __future__ import annotations + +from collections.abc import Mapping +from collections.abc import Sequence +from dataclasses import dataclass +from functools import partial +import inspect +import threading +from types import MappingProxyType +from typing import Any + +from cuda.tile import _backends as _backend +from cuda.tile._backends._signature import array_metadata, constant_key, constant_parameters +from cuda.tile.tune._tune import _benchmark_launch, exhaustive_search + + +@dataclass(frozen=True, init=False) +class Config: + """A set of compile-time values and compiler hints for autotuning.""" + + meta: Mapping[str, Any] + hints: Mapping[str, Any] + + def __init__(self, meta: Mapping[str, Any], /, **hints: Any): + if not isinstance(meta, Mapping): + raise TypeError("Config meta must be a mapping") + values = dict(meta) + if any(not isinstance(name, str) for name in values): + raise TypeError("Config meta names must be strings") + object.__setattr__(self, "meta", MappingProxyType(values)) + object.__setattr__(self, "hints", MappingProxyType(dict(hints))) + + def __repr__(self) -> str: + hints = "".join(f", {name}={value!r}" + for name, value in self.hints.items()) + return f"Config({dict(self.meta)!r}{hints})" + + +def _normalize_configs(configs, valid_names, what): + """Validate ``configs`` and return them with the set of tuned names.""" + + normalized = tuple(config if isinstance(config, Config) else Config(config) + for config in configs) + if not normalized: + raise ValueError("autotune configs must not be empty") + tuned = set(normalized[0].meta) + for config in normalized: + unknown = set(config.meta) - valid_names + if unknown: + raise ValueError( + f"Config contains unknown {what}: {sorted(unknown)}") + if set(config.meta) != tuned: + raise ValueError(f"All autotune configs must use the same {what}") + return normalized, tuned + + +def _key_names(key, valid_names, tuned): + names = tuple(key) + if len(set(names)) != len(names): + raise ValueError("autotune key names must be unique") + unknown = set(names) - valid_names + if unknown: + raise ValueError( + f"autotune key contains unknown parameters: {sorted(unknown)}") + overlap = set(names) & tuned + if overlap: + raise ValueError( + f"autotune key cannot contain autotuned parameters: {sorted(overlap)}") + return set(names) + + +def _array_key(value): + """Layout fingerprint of an array value, or ``None`` if not an array.""" + + try: + array = array_metadata(value) + except TypeError: + return None + return array.device, array.device_index, array.dtype, array.shape, array.strides + + +def _config_key(named_values, extra_names): + """Key deciding which tuning result applies to ``named_values``. + + Arrays contribute their layout only, so a tuning result is shared between + distinct buffers of the same shape. Other values contribute only when + named, either explicitly via ``key`` or because they are compile-time + constants. + """ + + key = [] + for name, value in named_values: + array = _array_key(value) + if array is not None: + key.append((name, array)) + elif name in extra_names: + key.append((name, constant_key(value))) + return tuple(key) + + +def _option_key(value): + if isinstance(value, Mapping): + return tuple(sorted((name, _option_key(item)) for name, item in value.items())) + if isinstance(value, (tuple, list)): + return tuple(map(_option_key, value)) + return constant_key(value) + + +def _environment_key(backend, options): + if backend is None: + return None + # Launch-only options affect tuning even when they do not affect compilation. + return backend, backend.normalize_options(options), _option_key(options) + + +def _make_injector(positions, values): + """Build a callable inserting ``values`` at ``positions`` (ascending).""" + + def inject(args): + full = list(args) + for position, value in zip(positions, values): + full.insert(position, value) + return tuple(full) + + return inject + + +class _Autotuned: + """Config cache shared by the two autotuning flavors.""" + + _configs: tuple[Config, ...] + + def __init__(self): + self._cache: dict[Any, int] = {} + # Serializes tuning, so concurrent first calls tune once. + self._lock = threading.Lock() + # The selected config and prepared launch are per thread. + self._local = threading.local() + + @property + def best_config(self) -> Config | None: + """The config selected by this thread's most recent call.""" + + best = getattr(self._local, "best", None) + return None if best is None else self._configs[best] + + def _index_by_id(self): + return {id(config): index for index, config in enumerate(self._configs)} + + def _select(self, key, tune): + """Return the config index for ``key``, tuning with ``tune()`` on a miss.""" + with self._lock: + index = self._cache.get(key) + result = None + if index is None: + index, result = tune() + self._cache[key] = index + self._local.best = index + return index, result + + def _prepared(self): + prepared = getattr(self._local, "prepared", None) + if prepared is None: + raise RuntimeError( + f"autotuned {self._kind} must be prepared before launch") + return prepared + + +class AutotunedKernel(_Autotuned): + """A lazily tuned wrapper around a :class:`cuda.tile.kernel`. + + Configs supply a subset of the kernel's ``ct.Constant`` parameters; those + parameters are omitted from the argument tuple passed at launch. + """ + + _kind = "kernel" + + def __init__(self, kernel, configs: Sequence[Config | Mapping[str, Any]], + *, key: Sequence[str] = (), grid, options=None, + quiet: bool = True): + super().__init__() + annotated = getattr(kernel, "_annotated_function", None) + if annotated is None: + raise TypeError("autotune must wrap a cuda.tile.kernel") + if grid is None: + raise TypeError("autotune requires a grid callable or tuple") + if options is not None and not callable(options): + if not isinstance(options, Mapping): + raise TypeError("autotune options must be a mapping or callable") + options = dict(options) + + names = tuple(annotated.pysig.parameters) + constants = {name for name, is_constant + in zip(names, constant_parameters(kernel)) + if is_constant} + self._configs, tuned = _normalize_configs(configs, constants, + "constant parameters") + # Untuned constants stay caller arguments but are still compile-time + # values, so they have to take part in config selection. + self._extra_names = (_key_names(key, set(names), tuned) + | (constants - tuned)) + positions = tuple(index for index, name in enumerate(names) + if name in tuned) + + self._kernel = kernel + self._names = names + self._arg_names = tuple(name for name in names if name not in tuned) + self._kernels = tuple( + kernel.replace_hints(**config.hints) if config.hints else kernel + for config in self._configs) + self._injectors = tuple( + _make_injector(positions, + tuple(config.meta[names[i]] for i in positions)) + for config in self._configs) + self._grid = grid + self._options = options + self._quiet = quiet + + @property + def kernel(self): + return self._kernel + + def _plan_for(self, index, args): + grid = self._grid + options = self._options + if callable(grid) or callable(options): + full_args = self._injectors[index](args) + meta = dict(zip(self._names, full_args)) + if callable(grid): + grid = grid(meta) + if callable(options): + options = options(meta) + if not isinstance(options, Mapping): + raise TypeError( + "autotune options callable must return a mapping") + options = dict(options) + return ((grid,) if isinstance(grid, int) else tuple(grid), options) + + @staticmethod + def _options_context_factory(options): + if options is None: + return None + return partial(_backend.backend_options, options) + + def _tune(self, stream, args): + index_by_id = self._index_by_id() + plans = {} + + def plan_for(config): + index = index_by_id[id(config)] + plan = plans.get(index) + if plan is None: + plan = self._plan_for(index, args) + plans[index] = plan + return plan + + def indexed_args(config): + index = index_by_id[id(config)] + return index, self._injectors[index](args) + + # Time the kernels that later launches use, so each config compiles once. + search_kwargs = {} + if self._options is not None: + search_kwargs["context_fn"] = lambda config: ( + self._options_context_factory(plan_for(config)[1])()) + result = exhaustive_search( + self._configs, + stream, + grid_fn=lambda config: plan_for(config)[0], + kernel=None, + args_fn=indexed_args, + quiet=self._quiet, + benchmark_fn=lambda stream, grid, _kernel, indexed, flush_l2: _benchmark_launch( + stream, grid, self._kernels[indexed[0]], indexed[1], flush_l2), + **search_kwargs, + ) + return index_by_id[id(result.best.config)], None + + def _resolve(self, stream, args): + if len(args) != len(self._arg_names): + raise TypeError( + f"autotuned kernel expects {len(self._arg_names)} arguments, " + f"got {len(args)}") + if len(self._configs) == 1: + self._local.best = 0 + return 0 + backend = _backend.get_backend() + if callable(self._options): + environment = tuple( + _environment_key(backend, self._plan_for(index, args)[1]) + for index in range(len(self._configs))) + else: + options = (_backend.current_options() if self._options is None + else self._options) + environment = _environment_key(backend, options) + key = environment, _config_key(zip(self._arg_names, args), self._extra_names) + index, _ = self._select(key, lambda: self._tune(stream, args)) + return index + + def prepare(self, stream, kernel_args, /) -> Config: + """Tune and compile for ``kernel_args`` ahead of timing. + + Use :meth:`launch_prepared` to benchmark without config lookup, + constant injection or grid computation in the measured call. + """ + + args = tuple(kernel_args) + index = self._resolve(stream, args) + grid, options = self._plan_for(index, args) + full_args = self._injectors[index](args) + kernel = self._kernels[index] + context_factory = self._options_context_factory(options) + if _backend.get_backend() is None: + launch_fn, target = _backend.launch, kernel + elif context_factory is None: + launch_fn = _backend.launch_compiled + target = _backend.compile_for_launch(kernel, full_args) + else: + launch_fn = _backend.launch_compiled + with context_factory(): + target = _backend.compile_for_launch(kernel, full_args) + self._local.prepared = (launch_fn, grid, target, full_args, context_factory) + return self._configs[index] + + def launch_prepared(self, stream, /): + """Launch what the most recent :meth:`prepare` call built. + + This path deliberately skips all lookups, so it stays bound to the + arguments that were passed to :meth:`prepare`. + """ + + launch_fn, grid, target, args, context_factory = self._prepared() + if context_factory is None: + return launch_fn(stream, grid, target, args) + with context_factory(): + return launch_fn(stream, grid, target, args) + + def __call__(self, stream, kernel_args, /): + """Launch with the best known config for ``kernel_args``.""" + + args = tuple(kernel_args) + index = self._resolve(stream, args) + grid, options = self._plan_for(index, args) + if options is None: + return _backend.launch(stream, grid, self._kernels[index], + self._injectors[index](args)) + with self._options_context_factory(options)(): + return _backend.launch(stream, grid, self._kernels[index], + self._injectors[index](args)) + + +class AutotunedFunction(_Autotuned): + """Autotune a host factory that returns a prepared callable. + + Each config overrides named factory parameters; the returned callable may + launch any number of kernels. + """ + + _kind = "function" + + def __init__(self, function, configs: Sequence[Config | Mapping[str, Any]], + *, key: Sequence[str] = (), quiet: bool = True): + super().__init__() + self._function = function + self._signature = inspect.signature(function) + names = set(self._signature.parameters) + self._configs, tuned = _normalize_configs(configs, names, + "function parameters") + if any(config.hints for config in self._configs): + raise ValueError( + "Compiler hints are only supported when autotuning a ct.kernel") + self._extra_names = _key_names(key, names, tuned) + self._quiet = quiet + + def _bind(self, args, kwargs): + bound = self._signature.bind(*args, **kwargs) + bound.apply_defaults() + return bound + + def _build(self, index, args, kwargs): + bound = self._bind(args, kwargs) + bound.arguments.update(self._configs[index].meta) + prepared = self._function(*bound.args, **bound.kwargs) + if not callable(prepared): + raise TypeError( + "autotuned function must return a callable prepared operation") + return prepared + + def _tune(self, args, kwargs): + index_by_id = self._index_by_id() + built: dict[int, Any] = {} + + def args_fn(config): + index = index_by_id[id(config)] + if index not in built: + built[index] = self._build(index, args, kwargs) + return (built[index],) + + result = exhaustive_search( + self._configs, + None, + grid_fn=lambda _config: (1,), + kernel=None, + args_fn=args_fn, + hints_fn=None, + quiet=self._quiet, + benchmark_fn=lambda stream, _grid, _kernel, bench_args, _flush_l2: ( + _backend.benchmark_callable(stream, bench_args[0])), + ) + index = index_by_id[id(result.best.config)] + return index, built[index] + + def _resolve(self, args, kwargs): + """Return the config index and the operation built while tuning, if any.""" + + key = (_environment_key(_backend.get_backend(), _backend.current_options()), + _config_key(self._bind(args, kwargs).arguments.items(), + self._extra_names)) + return self._select(key, lambda: self._tune(args, kwargs)) + + def prepare(self, *args, **kwargs) -> Config: + """Build the prepared multi-kernel operation ahead of timing.""" + + index, prepared = self._resolve(args, kwargs) + self._local.prepared = (prepared if prepared is not None + else self._build(index, args, kwargs)) + return self._configs[index] + + def launch_prepared(self): + """Run the operation built by the most recent :meth:`prepare` call.""" + + return self._prepared()() + + def __call__(self, *args, **kwargs): + index, prepared = self._resolve(args, kwargs) + if prepared is None: + prepared = self._build(index, args, kwargs) + return prepared() + + +def autotune(*, + configs: Sequence[Config | Mapping[str, Any]], + key: Sequence[str] = (), + grid=None, + options=None, + quiet: bool = True): + """Autotune a cuTile kernel or a factory for one or more kernel launches. + + ``configs`` is the search space. For a kernel, each config may provide a + subset of its ``ct.Constant`` parameters and compiler hints. Configured + constants are inserted into the kernel arguments automatically. For a + regular Python function, config entries override parameters by name and + the function must return a callable prepared operation. This supports + factories that prepare and capture several kernel launches. + + Kernel mode requires ``grid``. It can be a fixed integer or tuple, or a + callable with the form ``grid(meta)``. ``meta`` is a dictionary containing + every kernel parameter, including runtime arguments and configured + constants. This lets the grid depend on both input shapes and the config:: + + @ct.autotune( + configs=[ + ct.tune.Config({"TILE": 32}), + ct.tune.Config({"TILE": 64}), + ], + grid=lambda meta: ( + ct.cdiv(meta["x"].shape[0], meta["TILE"]), + ), + options=lambda meta: { + "assume_in_bounds": meta["x"].shape[0] % meta["TILE"] == 0, + }, + ) + @ct.kernel + def add_one(x, y, TILE: ct.Constant[int]): + bid = ct.bid(0) + value = ct.load(x, index=(bid,), shape=(TILE,)) + ct.store(y, index=(bid,), tile=value + 1) + + add_one.prepare(stream, (x, y)) + add_one.launch_prepared(stream) + + ``options`` accepts a fixed backend-options mapping or an + ``options(meta)`` callback. Like ``grid(meta)``, the callback receives the + complete parameter dictionary after configured constants are inserted and + must return an options mapping. Use it when a backend option depends on a + selected config, for example to enable CPU bounds assumptions only when a + tuned block size divides the input length:: + + @ct.autotune( + configs=[ + ct.tune.Config({"BLOCK_SIZE": 256}), + ct.tune.Config({"BLOCK_SIZE": 512}), + ], + grid=lambda meta: ( + ct.cdiv(meta["n_elements"], meta["BLOCK_SIZE"]), + ), + options=lambda meta: { + "assume_in_bounds": ( + meta["n_elements"] % meta["BLOCK_SIZE"] == 0 + ), + }, + ) + @ct.kernel + def leaky_relu(x, output, n_elements: ct.Constant[int], + BLOCK_SIZE: ct.Constant[int]): + ... + + The resolved options are active, as with :func:`cuda.tile.backend_options`, + while each candidate is tuned, and while the selected kernel is compiled + and launched. + + Factory mode omits ``grid`` because the returned prepared callable owns + the grids for its launches. The factory example below tunes a parameter + that changes a multi-kernel operation and keeps tuning, preparation, and + config lookup outside the measured region:: + + @ct.autotune( + configs=[ + ct.tune.Config({"blocking_factor_k": 1}), + ct.tune.Config({"blocking_factor_k": 2}), + ], + ) + def prepare_matmul(a, b, *, blocking_factor_k=1): + return prepare_sfc_matmul( + a, b, blocking_factor_k=blocking_factor_k) + + prepare_matmul.prepare(a, b) + result = prepare_matmul.launch_prepared() + + On the first call for a new input layout or key, every config is prepared + and benchmarked. Later calls reuse the selected config. Array arguments are + keyed by device, dtype, shape, and strides; use ``key`` for additional + scalar parameters that affect the generated work. Use ``prepare`` followed + by ``launch_prepared`` when measuring kernel time: the latter performs no + tuning, compilation, grid calculation, or argument preparation. The + ``stream`` argument is required in kernel mode and is omitted in factory + mode. + + Args: + configs: Non-empty sequence of :class:`Config` objects or metadata + mappings. All configs in one decorator must name the same + parameters. + key: Optional names of scalar runtime parameters that should select a + separate tuning result. Array layouts are included automatically. + grid: Fixed launch geometry or ``grid(meta)`` callable for kernel mode. + It is invalid for factory mode. + options: Fixed backend compile options mapping or ``options(meta)`` + callback for kernel mode. The callback receives configured + constants and runtime parameters and must return a mapping. The + selected options are active while tuning, compiling, and + launching. It is invalid for factory mode. + quiet: If true, suppress the tuning summary and progress output. + + Returns: + A callable wrapper with ``prepare`` and ``launch_prepared`` methods, + plus ``best_config``. Kernel wrappers are called as + ``wrapper(stream, kernel_args)``; factory wrappers are called with the + factory's normal arguments. + + Raises: + ValueError: If configs or key names are invalid. + TypeError: If the decorated object or grid form is invalid. + """ + + def decorate(kernel_or_function): + if getattr(kernel_or_function, "_annotated_function", None) is not None: + return AutotunedKernel(kernel_or_function, configs, key=key, + grid=grid, options=options, quiet=quiet) + if grid is not None: + raise TypeError("grid is only valid when autotuning a ct.kernel") + if options is not None: + raise TypeError("options is only valid when autotuning a ct.kernel") + return AutotunedFunction(kernel_or_function, configs, key=key, + quiet=quiet) + + return decorate + + +__all__ = ["AutotunedFunction", "AutotunedKernel", "Config", "autotune"] diff --git a/src/cuda/tile/tune/_tune.py b/src/cuda/tile/tune/_tune.py index 5ed69563..fd234076 100644 --- a/src/cuda/tile/tune/_tune.py +++ b/src/cuda/tile/tune/_tune.py @@ -5,9 +5,12 @@ import heapq import math +from contextlib import nullcontext from dataclasses import dataclass -from typing import Any, Callable, Generic, Sequence, TypeVar +from functools import partial +from typing import Any, Callable, ContextManager, Generic, Sequence, TypeVar +from cuda.tile import _backends as _backend from cuda.tile._cext import _benchmark from cuda.tile._execution import kernel as TileKernel from cuda.tile.tune._tune_utils import benchmark_with_timeout @@ -133,19 +136,21 @@ def exhaustive_search( search_space: Sequence[T], stream, grid_fn: Callable[[T], tuple[int, ...]], - kernel: TileKernel | Callable[[T], TileKernel], + kernel: TileKernel | Callable[[T], TileKernel] | None, args_fn: Callable[[T], tuple[Any, ...]], hints_fn: Callable[[T], dict[str, Any]] | None = None, *, quiet: bool = False, single_run_timeout_sec: float | None = None, flush_l2: bool = True, + benchmark_fn: Callable[..., float] | None = None, + context_fn: Callable[[T], ContextManager[Any]] | None = None, ) -> TuningResult[T]: """Searches the entire search space and return the best configuration. Args: search_space: Sequence of configs to evaluate. - stream: The CUDA stream to execute kernel on. + stream: The stream to execute kernel on (``None`` for the CPU backend). grid_fn: Maps a config to grid dimensions. kernel: The kernel to tune, or a function mapping a config to the kernel to tune. For best performance when using a function, return the same kernel @@ -156,9 +161,16 @@ def exhaustive_search( single_run_timeout_sec: Wall-time timeout (in seconds) per kernel launch, enforced by running benchmarks in a subprocess. When None (the default), timeouts are disabled and kernels run directly without a - subprocess. + subprocess. Requires CUDA; not supported by custom backends. flush_l2: Flush L2 before each kernel run. Defaults to True. Set to False - to measure without deliberately evicting cached data. + to measure without deliberately evicting cached data. Ignored by + custom backends. + benchmark_fn: Optional ``benchmark_fn(stream, grid, kernel, args, flush_l2)`` + timing one call in microseconds, used instead of launching ``kernel``. + ``kernel`` may then be ``None``. Not supported with + ``single_run_timeout_sec``. + context_fn: Optional context manager factory scoped around each + configuration's preparation and each of its timing runs. Returns: @@ -244,6 +256,10 @@ def tune(x, y, out) -> ct.tune.TuningResult: if total == 0: raise ValueError("Search space is empty.") + if kernel is None and benchmark_fn is None: + raise ValueError("kernel can only be None with benchmark_fn.") + if single_run_timeout_sec is not None and benchmark_fn is not None: + raise ValueError("single_run_timeout_sec is not supported with benchmark_fn.") # min-heap of running candidates, sorted by (mean_us, error_margin_us). # The integer tie-breaker keeps heap operations stable when timings match. @@ -256,16 +272,26 @@ def tune(x, y, out) -> ct.tune.TuningResult: for i, cfg in enumerate(search_space): if not quiet and isatty: progress(0, i, total, len(errors)) - grid = grid_fn(cfg) - base_kernel = kernel if isinstance(kernel, TileKernel) else kernel(cfg) - hints = hints_fn(cfg) if hints_fn is not None else {} - updated_kernel = base_kernel.replace_hints(**hints) + context = partial(context_fn, cfg) if context_fn is not None else nullcontext + with context(): + grid = grid_fn(cfg) + updated_kernel = None + if kernel is not None: + base_kernel = kernel if isinstance(kernel, TileKernel) else kernel(cfg) + hints = hints_fn(cfg) if hints_fn is not None else {} + updated_kernel = base_kernel.replace_hints(**hints) + if (single_run_timeout_sec is not None + and _backend.get_backend() is not None): + raise ValueError( + "single_run_timeout_sec is not supported by custom backends.") candidate = _TimingCandidate( config=cfg, grid=grid, kernel=updated_kernel, get_args=lambda _cfg=cfg: args_fn(_cfg), flush_l2=flush_l2, + benchmark_fn=benchmark_fn or _benchmark_launch, + context=context, ) try: @@ -348,6 +374,13 @@ def tune(x, y, out) -> ct.tune.TuningResult: _TOP_K = 5 +def _benchmark_launch(stream, grid, kernel, args, flush_l2: bool) -> float: + """Time one launch in microseconds where :func:`cuda.tile.launch` would run it.""" + if _backend.get_backend() is not None: + return _backend.benchmark(stream, grid, kernel, args) + return _benchmark(stream, grid, kernel, args, flush_l2) + + @dataclass class _TimingCandidate(Generic[T]): config: T @@ -355,6 +388,8 @@ class _TimingCandidate(Generic[T]): kernel: Any get_args: Callable[[], tuple[Any, ...]] flush_l2: bool + benchmark_fn: Callable[..., float] + context: Callable[[], ContextManager[Any]] num_samples: int = 0 mean_us: float = 0.0 m2: float = 0.0 @@ -363,19 +398,24 @@ class _TimingCandidate(Generic[T]): def warmup(self, stream, num_times, launch_timeout_sec): assert num_times > 0 - for i in range(num_times): - if i == 0 and launch_timeout_sec is not None: - # First warmup is timed to ensure it doesn't deadlock. - benchmark_with_timeout( - stream, self.grid, self.kernel, self.get_args(), launch_timeout_sec, - self.flush_l2) - else: - _benchmark(stream, self.grid, self.kernel, self.get_args(), self.flush_l2) + with self.context(): + for i in range(num_times): + if i == 0 and launch_timeout_sec is not None: + # First warmup is timed to ensure it doesn't deadlock. + benchmark_with_timeout( + stream, self.grid, self.kernel, self.get_args(), launch_timeout_sec, + self.flush_l2) + else: + self._time(stream) def run_benchmark(self, stream, num_times): - for _ in range(min(num_times, _MAX_REPEATS - self.num_samples)): - self._add_sample(_benchmark( - stream, self.grid, self.kernel, self.get_args(), self.flush_l2)) + with self.context(): + for _ in range(min(num_times, _MAX_REPEATS - self.num_samples)): + self._add_sample(self._time(stream)) + + def _time(self, stream) -> float: + return self.benchmark_fn( + stream, self.grid, self.kernel, self.get_args(), self.flush_l2) def converged(self) -> bool: # Stop if ... diff --git a/test/test_autotune.py b/test/test_autotune.py new file mode 100644 index 00000000..50440df3 --- /dev/null +++ b/test/test_autotune.py @@ -0,0 +1,313 @@ +# SPDX-FileCopyrightText: Copyright (c) <2026> Intel Corporation. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch + +import cuda.tile as ct +import cuda.tile.tune._autotune as autotune_mod +from cuda.tile.tune import Config, autotune + + +@ct.kernel +def _kernel(x, N: ct.Constant[int], TILE: ct.Constant[int], FLAG: ct.Constant[bool]): + pass + + +@ct.kernel +def _runtime_kernel(x, scale: float, TILE: ct.Constant[int]): + pass + + +def _fake_runtime(monkeypatch, events=None): + tuned = [] + launched = [] + active_options = [] + + def search(search_space, stream, *, grid_fn, kernel, args_fn, quiet, + hints_fn=None, benchmark_fn=None, context_fn=None): + config = search_space[0] + if context_fn is None: + tuned.append((config, args_fn(config), grid_fn(config))) + else: + with context_fn(config): + tuned.append((config, args_fn(config), grid_fn(config))) + return SimpleNamespace(best=SimpleNamespace(config=config)) + + def compile_kernel(kernel, args): + if events is not None and active_options: + events.append(("compile", dict(active_options[-1]))) + return SimpleNamespace(binary=object(), symbol="kernel", signature=None) + + def launch_compiled(stream, grid, compiled, args): + if events is not None and active_options: + events.append(("launch", dict(active_options[-1]))) + launched.append((grid, args)) + + def launch(stream, grid, kernel, args): + if events is not None and active_options: + events.append(("launch", dict(active_options[-1]))) + launched.append((grid, args)) + + @contextmanager + def backend_options(options): + options = dict(options) + if events is not None: + events.append(("enter", options)) + active_options.append(options) + try: + yield + finally: + active_options.pop() + if events is not None: + events.append(("exit", options)) + + monkeypatch.setattr(autotune_mod, "exhaustive_search", search) + backend = autotune_mod._backend + + class Backend: + def normalize_options(self, options): + return () + + target = Backend() + monkeypatch.setattr(backend, "get_backend", lambda: target) + monkeypatch.setattr(backend, "compile_for_launch", compile_kernel) + monkeypatch.setattr(backend, "launch_compiled", launch_compiled) + monkeypatch.setattr(backend, "backend_options", backend_options) + monkeypatch.setattr(backend, "launch", launch) + return tuned, launched + + +def test_autotune_injects_configured_constants_and_reuses_tuning(monkeypatch): + tuned, launched = _fake_runtime(monkeypatch) + autotuned = autotune( + configs=[Config({"TILE": 4}), Config({"TILE": 8})], + grid=lambda meta: (meta["x"].shape[0] // meta["TILE"],), + )(_kernel) + + x = torch.empty(16) + args = (x, 16, True) + autotuned.prepare(None, args) + autotuned(None, args) + autotuned(None, args) + + assert len(tuned) == 1 + assert len(launched) == 2 + assert launched[0] == ((4,), (x, 16, 4, True)) + assert autotuned.best_config == tuned[0][0] + + +def test_launch_prepared_uses_compiled_kernel(monkeypatch): + _tuned, launched = _fake_runtime(monkeypatch) + autotuned = autotune( + configs=[Config({"TILE": 4})], + grid=(1,), + )(_runtime_kernel) + args = (torch.empty(16), 1.0) + + with pytest.raises(RuntimeError, match="prepared"): + autotuned.launch_prepared(None) + autotuned.prepare(None, args) + autotuned.launch_prepared(None) + + assert launched == [((1,), (args[0], 1.0, 4))] + + +def test_autotune_options_scope_tuning_compile_and_launch(monkeypatch): + events = [] + tuned, launched = _fake_runtime(monkeypatch, events) + options = {"assume_in_bounds": True, "num_cpu_threads": 4} + autotuned = autotune( + configs=[Config({"TILE": 4})], + grid=(1,), + options=options, + )(_runtime_kernel) + args = (torch.empty(16), 1.0) + + autotuned.prepare(None, args) + autotuned.launch_prepared(None) + + assert launched == [((1,), (args[0], 1.0, 4))] + assert tuned == [] + assert events == [ + ("enter", options), ("compile", options), ("exit", options), + ("enter", options), ("launch", options), ("exit", options), + ] + + +def test_autotune_options_callable_receives_configured_meta(monkeypatch): + tuned, _launched = _fake_runtime(monkeypatch) + received = [] + + def options(meta): + received.append(meta) + return {"assume_in_bounds": meta["N"] % meta["TILE"] == 0} + + autotuned = autotune( + configs=[Config({"TILE": 4})], + grid=(1,), + options=options, + )(_kernel) + x = torch.empty(16) + autotuned(None, (x, 16, True)) + + assert tuned == [] + assert len(received) == 1 + assert all(meta["x"] is x and meta["N"] == 16 + and meta["TILE"] == 4 and meta["FLAG"] is True + for meta in received) + + +def test_autotune_rejects_invalid_options_callable_result(monkeypatch): + _fake_runtime(monkeypatch) + autotuned = autotune( + configs=[Config({"TILE": 4})], + grid=(1,), + options=lambda _meta: None, + )(_runtime_kernel) + + with pytest.raises(TypeError, match="options callable must return a mapping"): + autotuned(None, (torch.empty(16), 1.0)) + + +def test_single_config_skips_tuning_for_different_keys(monkeypatch): + tuned, launched = _fake_runtime(monkeypatch) + autotuned = autotune( + configs=[Config({"TILE": 4})], + grid=lambda meta: (meta["x"].shape[0] // meta["TILE"],), + )(_kernel) + + autotuned(None, (torch.empty(16), 16, True)) + autotuned(None, (torch.empty(16), 32, True)) + + assert tuned == [] + assert [entry[1][1] for entry in launched] == [16, 32] + + +def test_unkeyed_runtime_scalar_cannot_reuse_stale_grid(monkeypatch): + tuned, launched = _fake_runtime(monkeypatch) + autotuned = autotune( + configs=[Config({"TILE": 4})], + grid=lambda meta: (int(meta["scale"]),), + )(_runtime_kernel) + + x = torch.empty(16) + autotuned(None, (x, 1.0)) + autotuned(None, (x, 2.0)) + + assert tuned == [] + assert [entry[0] for entry in launched] == [(1,), (2,)] + + +def test_autotune_rejects_inconsistent_configured_constants(): + with pytest.raises(ValueError, match="same constant parameters"): + autotune( + configs=[Config({"TILE": 4}), Config({"N": 16})], + grid=(1,), + )(_kernel) + + +def test_autotune_rejects_unknown_constant_parameter(): + with pytest.raises(ValueError, match="unknown constant parameters"): + autotune(configs=[Config({"UNKNOWN": 4})], grid=(1,))(_kernel) + + +def test_autotune_forwards_generic_compiler_hints(monkeypatch): + received = [] + + def replace_hints(kernel, **hints): + received.append(hints) + return kernel + + monkeypatch.setattr( + type(_runtime_kernel), "replace_hints", replace_hints) + autotune( + configs=[Config({"TILE": 4}, future_compiler_hint="value")], + grid=(1,), + )(_runtime_kernel) + + assert received == [{"future_compiler_hint": "value"}] + + +def test_autotune_rejects_options_for_factory(): + with pytest.raises(TypeError, match="options is only valid"): + autotune(configs=[Config({})], options={})(lambda: lambda: None) + + +def test_autotune_key_preserves_float_constant_bits(): + assert autotune_mod._config_key([("value", 0.0)], {"value"}) != ( + autotune_mod._config_key([("value", -0.0)], {"value"})) + + +def test_autotune_factory_reuses_prepared_multi_kernel_operation(monkeypatch): + prepared_calls = [] + factory_calls = [] + timer_calls = [] + + def factory(x, blocking_factor_k=1): + factory_calls.append(blocking_factor_k) + + def prepared(): + prepared_calls.append(blocking_factor_k) + + return prepared + + def timer(stream, prepared, args=()): + timer_calls.append(prepared) + prepared(*args) + return 1.0 + + def search(search_space, stream, *, grid_fn, kernel, args_fn, + hints_fn, quiet, benchmark_fn): + config = search_space[0] + benchmark_fn(stream, grid_fn(config), kernel, args_fn(config), True) + return SimpleNamespace(best=SimpleNamespace(config=config)) + + monkeypatch.setattr(autotune_mod, "exhaustive_search", search) + monkeypatch.setattr(autotune_mod._backend, "benchmark_callable", timer) + + autotuned = autotune( + configs=[ + Config({"blocking_factor_k": 1}), + Config({"blocking_factor_k": 2}), + ], + )(factory) + x = torch.empty(16) + autotuned.prepare(x) + autotuned.launch_prepared() + autotuned.launch_prepared() + + assert factory_calls == [1] + assert len(timer_calls) == 1 + assert prepared_calls == [1, 1, 1] + + +def test_factory_rebuilds_operation_for_different_storage(monkeypatch): + factory_calls = [] + + def factory(x): + factory_calls.append(x) + return lambda: None + + def search(search_space, stream, *, args_fn, **kwargs): + config = search_space[0] + args_fn(config) + return SimpleNamespace(best=SimpleNamespace(config=config)) + + monkeypatch.setattr(autotune_mod, "exhaustive_search", search) + monkeypatch.setattr( + autotune_mod._backend, + "benchmark_callable", + lambda stream, fn, args=(): 1.0, + ) + autotuned = autotune(configs=[Config({})])(factory) + first = torch.empty(16) + second = torch.empty(16) + autotuned.prepare(first) + autotuned.prepare(second) + + assert factory_calls == [first, second] diff --git a/uv.lock b/uv.lock index c809ddbb..c8d4d231 100644 --- a/uv.lock +++ b/uv.lock @@ -1,6 +1,38 @@ version = 1 revision = 3 requires-python = ">=3.11, <3.16" +resolution-markers = [ + "python_full_version >= '3.12' and sys_platform != 'darwin'", + "python_full_version >= '3.12' and sys_platform == 'darwin'", + "python_full_version < '3.12' and sys_platform != 'darwin'", + "python_full_version < '3.12' and sys_platform == 'darwin'", +] + +[[package]] +name = "cmake" +version = "3.31.10" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/37/7b/fbadb3f4fe90ad6ef57f9f5f9e4f721af8e86376fbdf11da2c6ed099830e/cmake-3.31.10.tar.gz", hash = "sha256:ec3d14a0e72e401b3665034dc37901df17f0b4e9c5b163be6cfedfb93470ac0f", size = 34499, upload-time = "2025-11-20T17:07:54.664Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/99/3b/6ed408a99709808df4014bead522e075f0e8d1100e6d886eb5bc30fee04e/cmake-3.31.10-py3-none-macosx_10_10_universal2.whl", hash = "sha256:ad697643a00d9ba85179590a383c4f7401169b55ebf4b8b2938daf28c6bdeb6d", size = 48001731, upload-time = "2025-11-20T17:06:57.473Z" }, + { url = "https://files.pythonhosted.org/packages/ad/59/acc2f79180d8aaf55b707ee65b5296ae76a3295ebff9e65d840849b6abc0/cmake-3.31.10-py3-none-manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:7c300e4ae68fbc1414a85505f7feb262cee82ff3304a286885ebf803b11a997c", size = 27579802, upload-time = "2025-11-20T17:07:00.676Z" }, + { url = "https://files.pythonhosted.org/packages/cc/0a/2f17b5cc0d3ac1ce2324364c044d170d06167f8d6b7464ee76857114d95e/cmake-3.31.10-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:3c17bb24dba15f8ecc3fd706afe04264410ef88796f4115c119327c961d5dc57", size = 26832178, upload-time = "2025-11-20T17:07:03.475Z" }, + { url = "https://files.pythonhosted.org/packages/b1/6c/d629c3b6a30105f5946e754b42d6ce84be36c190cc70b70fb2eb8a0ffc37/cmake-3.31.10-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a4a5615f31f692c9b9aa8b365704e4b76172348af6fa40e16fea3f118bb01194", size = 27165148, upload-time = "2025-11-20T17:07:06.895Z" }, + { url = "https://files.pythonhosted.org/packages/9b/d8/9842811ec5615598003499e2d38062d42651a20d022f93f9a590807d76b3/cmake-3.31.10-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:8d2ec5fc45d305227020c82213140a51a0cebe3c84f0299036f05716b3a52f60", size = 28889720, upload-time = "2025-11-20T17:07:09.597Z" }, + { url = "https://files.pythonhosted.org/packages/01/e9/7042b018121ceaae48416d37edb9924646e3cc8bbb417374cafdf8c2fa58/cmake-3.31.10-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a7864605238759a4ae8e3bd1fda2bb03978e3e37df310852662dcf53866413c8", size = 30757069, upload-time = "2025-11-20T17:07:12.537Z" }, + { url = "https://files.pythonhosted.org/packages/24/a8/4e320cd6dfae630c5d9532d822c20a095565bcb159be839919ff7ea2952f/cmake-3.31.10-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:1debc3a0823ce5d8d1bc17154599bbbb337c2681f93622b618bc78f46576e42a", size = 26932214, upload-time = "2025-11-20T17:07:16.075Z" }, + { url = "https://files.pythonhosted.org/packages/f4/2e/cefc60143950ca1a8dda26d4c8484e6fc406db16e7ca098ebae549de35f9/cmake-3.31.10-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2f766bb46367e5e0559fa33184653754bce044583a06014dcaebf8e6dff8a1f1", size = 27808794, upload-time = "2025-11-20T17:07:18.974Z" }, + { url = "https://files.pythonhosted.org/packages/1f/55/3f7d68d03f116142ce1076d65898e3d31c24d72985e827bfd9c601c8bc65/cmake-3.31.10-py3-none-manylinux_2_31_armv7l.whl", hash = "sha256:91410816db3beefe2f6032d721f9978c98dc7646e9992c0325486597164fab81", size = 24986234, upload-time = "2025-11-20T17:07:21.596Z" }, + { url = "https://files.pythonhosted.org/packages/e1/2e/634e46413472be742f3f422f159c41079bdc8ffba67538228f94a166c9ee/cmake-3.31.10-py3-none-musllinux_1_1_aarch64.whl", hash = "sha256:c2e5361dea9754ed3b06cf834894fb47dcbe7036d5e5d87acaeb10ff3dd5fd10", size = 27849273, upload-time = "2025-11-20T17:07:24.369Z" }, + { url = "https://files.pythonhosted.org/packages/6f/af/fac7ee4d79b688b94b4f6a6e5c4b080dca7594342576ad88a52b02a2d4eb/cmake-3.31.10-py3-none-musllinux_1_1_i686.whl", hash = "sha256:6970bb75c4dfc28cc31ff0cd848194d09094ae00d605181e1345b2ff70b61050", size = 31391299, upload-time = "2025-11-20T17:07:27.228Z" }, + { url = "https://files.pythonhosted.org/packages/d4/14/62339280a675106227ff650166c2899fe9f822ef9b855365563e71115125/cmake-3.31.10-py3-none-musllinux_1_1_ppc64le.whl", hash = "sha256:4cefb0a28ac1268b4eed4b595bf3aaff8de9704089066027700ec36584eccab8", size = 32105563, upload-time = "2025-11-20T17:07:34.352Z" }, + { url = "https://files.pythonhosted.org/packages/ba/19/a3014beb02b344599e7a879f9c91d5e388aa76924b1369e13ffd535d613d/cmake-3.31.10-py3-none-musllinux_1_1_s390x.whl", hash = "sha256:b331984de38dbda22d676f8812c8905526341ba7b397fe8c359255ff4d051193", size = 27972718, upload-time = "2025-11-20T17:07:37.492Z" }, + { url = "https://files.pythonhosted.org/packages/b0/53/63b70d583c93352beeb692f0ec6bf4dede96f92df85c009577a663be304a/cmake-3.31.10-py3-none-musllinux_1_1_x86_64.whl", hash = "sha256:678fc23db37cc69f01e18eb28790450ecc9401fd2fcd43364cc18f92330c12c2", size = 29497491, upload-time = "2025-11-20T17:07:40.316Z" }, + { url = "https://files.pythonhosted.org/packages/fc/a3/9d3a3881e70b947abe2d779e4c37d1aa9ac7f8339354e78d6036d520a220/cmake-3.31.10-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:e135d5f1e59f1dc1c80eea87321977d6c0eff86cb601c28e0965a1dd457ea587", size = 33183971, upload-time = "2025-11-20T17:07:43.392Z" }, + { url = "https://files.pythonhosted.org/packages/06/f2/8e13be79373ef808659e52680643e0388057060e2627757b90df72c50681/cmake-3.31.10-py3-none-win32.whl", hash = "sha256:b059a1810a2ce766b3e531bdc8d730bc192e260a9fa7dec7a0eb7a053d6063c7", size = 33416497, upload-time = "2025-11-20T17:07:46.467Z" }, + { url = "https://files.pythonhosted.org/packages/84/4b/e433ab430c580ec8e3d54cb2ec5c7ea76ae0b1e6ef2c2998c854f21d7780/cmake-3.31.10-py3-none-win_amd64.whl", hash = "sha256:f1ea1fe826355560e8976c3d5794d9357444209bc0e0d56676c71e6a571fd474", size = 36630429, upload-time = "2025-11-20T17:07:49.445Z" }, + { url = "https://files.pythonhosted.org/packages/9f/c3/4eb3288ccd779556fee4b7b0955bdb14545bbae83f85f7c819dde8013708/cmake-3.31.10-py3-none-win_arm64.whl", hash = "sha256:422a54711aa977af19d59b8f6010354cdda0b72a2e6d702b6d892e3e2cdf98a2", size = 35457178, upload-time = "2025-11-20T17:07:52.367Z" }, +] [[package]] name = "cuda-tile" @@ -9,14 +41,607 @@ dependencies = [ { name = "typing-extensions" }, ] +[package.optional-dependencies] +tileiras = [ + { name = "cuda-toolkit", extra = ["nvcc", "nvvm", "tileiras"] }, +] + +[package.dev-dependencies] +backend-build = [ + { name = "cmake" }, + { name = "cuda-toolkit", extra = ["cudart", "nvvm"], marker = "sys_platform == 'linux'" }, + { name = "nanobind" }, + { name = "ninja" }, + { name = "pybind11" }, + { name = "setuptools" }, + { name = "wheel" }, +] +cpu = [ + { name = "cmake" }, + { name = "cuda-toolkit", extra = ["cudart", "nvvm"], marker = "sys_platform == 'linux'" }, + { name = "nanobind" }, + { name = "ninja" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "numpy", version = "2.5.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "pybind11" }, + { name = "setuptools" }, + { name = "torch", version = "2.13.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform == 'darwin'" }, + { name = "torch", version = "2.13.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform != 'darwin'" }, + { name = "wheel" }, +] + [package.metadata] -requires-dist = [{ name = "typing-extensions" }] +requires-dist = [ + { name = "cuda-toolkit", extras = ["nvcc", "nvvm", "tileiras"], marker = "extra == 'tileiras'", specifier = ">=13.2,<13.6" }, + { name = "typing-extensions" }, +] +provides-extras = ["tileiras"] + +[package.metadata.requires-dev] +backend-build = [ + { name = "cmake", specifier = ">=3.24,<4" }, + { name = "cuda-toolkit", extras = ["cudart", "nvvm"], marker = "sys_platform == 'linux'", specifier = ">=13.2,<13.6" }, + { name = "nanobind", specifier = ">=2.9,<3" }, + { name = "ninja" }, + { name = "pybind11", specifier = ">=2.13.1" }, + { name = "setuptools", specifier = "==84.0.0" }, + { name = "wheel" }, +] +cpu = [ + { name = "cmake", specifier = ">=3.24,<4" }, + { name = "cuda-toolkit", extras = ["cudart", "nvvm"], marker = "sys_platform == 'linux'", specifier = ">=13.2,<13.6" }, + { name = "nanobind", specifier = ">=2.9,<3" }, + { name = "ninja" }, + { name = "numpy" }, + { name = "pybind11", specifier = ">=2.13.1" }, + { name = "setuptools", specifier = "==84.0.0" }, + { name = "torch", index = "https://download.pytorch.org/whl/cpu", conflict = { package = "cuda-tile", group = "cpu" } }, + { name = "wheel" }, +] + +[[package]] +name = "cuda-toolkit" +version = "13.3.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/61/a1/54c1e9498ba0df91ca15a46f41af6320cb9faed6ec2dbb30b6cbff8887c4/cuda_toolkit-13.3.1-py2.py3-none-any.whl", hash = "sha256:2ceda460a540323d52469bcfde48b48c1861f6482e4b5ea3cb5bdac00a1b11bd", size = 2656, upload-time = "2026-06-29T17:23:23.848Z" }, +] + +[package.optional-dependencies] +cudart = [ + { name = "nvidia-cuda-runtime", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, +] +nvcc = [ + { name = "nvidia-cuda-crt", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "nvidia-cuda-nvcc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "nvidia-cuda-runtime", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "nvidia-nvvm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, +] +nvvm = [ + { name = "nvidia-nvvm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, +] +tileiras = [ + { name = "nvidia-cuda-nvcc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "nvidia-cuda-tileiras", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "nvidia-nvvm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, +] + +[[package]] +name = "filelock" +version = "3.32.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7d/64/a02e6765de08964ed371eca577870593245afc9dfac16d037de7c10d18e6/filelock-3.32.3.tar.gz", hash = "sha256:0ffa185a3540854c95caa7fa76b76cb219d907415e2c5dc9af25fd970563487f", size = 218135, upload-time = "2026-08-13T16:00:05.577Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/8e/50f46a9c0ce8d2861a394c1347caae037ea0431d2f67d7feb151cbc4649a/filelock-3.32.3-py3-none-any.whl", hash = "sha256:7f0ca4bcc0e181c60dbbd8aa9ab5b120ebb99e4e064e83636340056f833a1f09", size = 98901, upload-time = "2026-08-13T16:00:03.974Z" }, +] + +[[package]] +name = "fsspec" +version = "2026.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/00/78/f34251dadb8f3921264a1d9b8946f5e542014ee2614b285261b4e40e6775/fsspec-2026.7.0.tar.gz", hash = "sha256:c803c40f4cf860b49dea58ee3e1c33cb9c790520e233537e1340049f89b82a88", size = 317040, upload-time = "2026-07-28T16:34:51.052Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl", hash = "sha256:b57ddbafedfaef7018c1ecab32aa200a9d7ca26b77965f64e48b70061249d279", size = 206583, upload-time = "2026-07-28T16:34:49.538Z" }, +] + +[[package]] +name = "jinja2" +version = "3.1.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, +] + +[[package]] +name = "markupsafe" +version = "3.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/08/db/fefacb2136439fc8dd20e797950e749aa1f4997ed584c62cfb8ef7c2be0e/markupsafe-3.0.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:1cc7ea17a6824959616c525620e387f6dd30fec8cb44f649e31712db02123dad", size = 11631, upload-time = "2025-09-27T18:36:18.185Z" }, + { url = "https://files.pythonhosted.org/packages/e1/2e/5898933336b61975ce9dc04decbc0a7f2fee78c30353c5efba7f2d6ff27a/markupsafe-3.0.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:4bd4cd07944443f5a265608cc6aab442e4f74dff8088b0dfc8238647b8f6ae9a", size = 12058, upload-time = "2025-09-27T18:36:19.444Z" }, + { url = "https://files.pythonhosted.org/packages/1d/09/adf2df3699d87d1d8184038df46a9c80d78c0148492323f4693df54e17bb/markupsafe-3.0.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6b5420a1d9450023228968e7e6a9ce57f65d148ab56d2313fcd589eee96a7a50", size = 24287, upload-time = "2025-09-27T18:36:20.768Z" }, + { url = "https://files.pythonhosted.org/packages/30/ac/0273f6fcb5f42e314c6d8cd99effae6a5354604d461b8d392b5ec9530a54/markupsafe-3.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0bf2a864d67e76e5c9a34dc26ec616a66b9888e25e7b9460e1c76d3293bd9dbf", size = 22940, upload-time = "2025-09-27T18:36:22.249Z" }, + { url = "https://files.pythonhosted.org/packages/19/ae/31c1be199ef767124c042c6c3e904da327a2f7f0cd63a0337e1eca2967a8/markupsafe-3.0.3-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:bc51efed119bc9cfdf792cdeaa4d67e8f6fcccab66ed4bfdd6bde3e59bfcbb2f", size = 21887, upload-time = "2025-09-27T18:36:23.535Z" }, + { url = "https://files.pythonhosted.org/packages/b2/76/7edcab99d5349a4532a459e1fe64f0b0467a3365056ae550d3bcf3f79e1e/markupsafe-3.0.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:068f375c472b3e7acbe2d5318dea141359e6900156b5b2ba06a30b169086b91a", size = 23692, upload-time = "2025-09-27T18:36:24.823Z" }, + { url = "https://files.pythonhosted.org/packages/a4/28/6e74cdd26d7514849143d69f0bf2399f929c37dc2b31e6829fd2045b2765/markupsafe-3.0.3-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:7be7b61bb172e1ed687f1754f8e7484f1c8019780f6f6b0786e76bb01c2ae115", size = 21471, upload-time = "2025-09-27T18:36:25.95Z" }, + { url = "https://files.pythonhosted.org/packages/62/7e/a145f36a5c2945673e590850a6f8014318d5577ed7e5920a4b3448e0865d/markupsafe-3.0.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:f9e130248f4462aaa8e2552d547f36ddadbeaa573879158d721bbd33dfe4743a", size = 22923, upload-time = "2025-09-27T18:36:27.109Z" }, + { url = "https://files.pythonhosted.org/packages/0f/62/d9c46a7f5c9adbeeeda52f5b8d802e1094e9717705a645efc71b0913a0a8/markupsafe-3.0.3-cp311-cp311-win32.whl", hash = "sha256:0db14f5dafddbb6d9208827849fad01f1a2609380add406671a26386cdf15a19", size = 14572, upload-time = "2025-09-27T18:36:28.045Z" }, + { url = "https://files.pythonhosted.org/packages/83/8a/4414c03d3f891739326e1783338e48fb49781cc915b2e0ee052aa490d586/markupsafe-3.0.3-cp311-cp311-win_amd64.whl", hash = "sha256:de8a88e63464af587c950061a5e6a67d3632e36df62b986892331d4620a35c01", size = 15077, upload-time = "2025-09-27T18:36:29.025Z" }, + { url = "https://files.pythonhosted.org/packages/35/73/893072b42e6862f319b5207adc9ae06070f095b358655f077f69a35601f0/markupsafe-3.0.3-cp311-cp311-win_arm64.whl", hash = "sha256:3b562dd9e9ea93f13d53989d23a7e775fdfd1066c33494ff43f5418bc8c58a5c", size = 13876, upload-time = "2025-09-27T18:36:29.954Z" }, + { url = "https://files.pythonhosted.org/packages/5a/72/147da192e38635ada20e0a2e1a51cf8823d2119ce8883f7053879c2199b5/markupsafe-3.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d53197da72cc091b024dd97249dfc7794d6a56530370992a5e1a08983ad9230e", size = 11615, upload-time = "2025-09-27T18:36:30.854Z" }, + { url = "https://files.pythonhosted.org/packages/9a/81/7e4e08678a1f98521201c3079f77db69fb552acd56067661f8c2f534a718/markupsafe-3.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1872df69a4de6aead3491198eaf13810b565bdbeec3ae2dc8780f14458ec73ce", size = 12020, upload-time = "2025-09-27T18:36:31.971Z" }, + { url = "https://files.pythonhosted.org/packages/1e/2c/799f4742efc39633a1b54a92eec4082e4f815314869865d876824c257c1e/markupsafe-3.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a7e8ae81ae39e62a41ec302f972ba6ae23a5c5396c8e60113e9066ef893da0d", size = 24332, upload-time = "2025-09-27T18:36:32.813Z" }, + { url = "https://files.pythonhosted.org/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d", size = 22947, upload-time = "2025-09-27T18:36:33.86Z" }, + { url = "https://files.pythonhosted.org/packages/2c/54/887f3092a85238093a0b2154bd629c89444f395618842e8b0c41783898ea/markupsafe-3.0.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:94c6f0bb423f739146aec64595853541634bde58b2135f27f61c1ffd1cd4d16a", size = 21962, upload-time = "2025-09-27T18:36:35.099Z" }, + { url = "https://files.pythonhosted.org/packages/c9/2f/336b8c7b6f4a4d95e91119dc8521402461b74a485558d8f238a68312f11c/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:be8813b57049a7dc738189df53d69395eba14fb99345e0a5994914a3864c8a4b", size = 23760, upload-time = "2025-09-27T18:36:36.001Z" }, + { url = "https://files.pythonhosted.org/packages/32/43/67935f2b7e4982ffb50a4d169b724d74b62a3964bc1a9a527f5ac4f1ee2b/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:83891d0e9fb81a825d9a6d61e3f07550ca70a076484292a70fde82c4b807286f", size = 21529, upload-time = "2025-09-27T18:36:36.906Z" }, + { url = "https://files.pythonhosted.org/packages/89/e0/4486f11e51bbba8b0c041098859e869e304d1c261e59244baa3d295d47b7/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b", size = 23015, upload-time = "2025-09-27T18:36:37.868Z" }, + { url = "https://files.pythonhosted.org/packages/2f/e1/78ee7a023dac597a5825441ebd17170785a9dab23de95d2c7508ade94e0e/markupsafe-3.0.3-cp312-cp312-win32.whl", hash = "sha256:d88b440e37a16e651bda4c7c2b930eb586fd15ca7406cb39e211fcff3bf3017d", size = 14540, upload-time = "2025-09-27T18:36:38.761Z" }, + { url = "https://files.pythonhosted.org/packages/aa/5b/bec5aa9bbbb2c946ca2733ef9c4ca91c91b6a24580193e891b5f7dbe8e1e/markupsafe-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:26a5784ded40c9e318cfc2bdb30fe164bdb8665ded9cd64d500a34fb42067b1c", size = 15105, upload-time = "2025-09-27T18:36:39.701Z" }, + { url = "https://files.pythonhosted.org/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f", size = 13906, upload-time = "2025-09-27T18:36:40.689Z" }, + { url = "https://files.pythonhosted.org/packages/38/2f/907b9c7bbba283e68f20259574b13d005c121a0fa4c175f9bed27c4597ff/markupsafe-3.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e1cf1972137e83c5d4c136c43ced9ac51d0e124706ee1c8aa8532c1287fa8795", size = 11622, upload-time = "2025-09-27T18:36:41.777Z" }, + { url = "https://files.pythonhosted.org/packages/9c/d9/5f7756922cdd676869eca1c4e3c0cd0df60ed30199ffd775e319089cb3ed/markupsafe-3.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:116bb52f642a37c115f517494ea5feb03889e04df47eeff5b130b1808ce7c219", size = 12029, upload-time = "2025-09-27T18:36:43.257Z" }, + { url = "https://files.pythonhosted.org/packages/00/07/575a68c754943058c78f30db02ee03a64b3c638586fba6a6dd56830b30a3/markupsafe-3.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:133a43e73a802c5562be9bbcd03d090aa5a1fe899db609c29e8c8d815c5f6de6", size = 24374, upload-time = "2025-09-27T18:36:44.508Z" }, + { url = "https://files.pythonhosted.org/packages/a9/21/9b05698b46f218fc0e118e1f8168395c65c8a2c750ae2bab54fc4bd4e0e8/markupsafe-3.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ccfcd093f13f0f0b7fdd0f198b90053bf7b2f02a3927a30e63f3ccc9df56b676", size = 22980, upload-time = "2025-09-27T18:36:45.385Z" }, + { url = "https://files.pythonhosted.org/packages/7f/71/544260864f893f18b6827315b988c146b559391e6e7e8f7252839b1b846a/markupsafe-3.0.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:509fa21c6deb7a7a273d629cf5ec029bc209d1a51178615ddf718f5918992ab9", size = 21990, upload-time = "2025-09-27T18:36:46.916Z" }, + { url = "https://files.pythonhosted.org/packages/c2/28/b50fc2f74d1ad761af2f5dcce7492648b983d00a65b8c0e0cb457c82ebbe/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a4afe79fb3de0b7097d81da19090f4df4f8d3a2b3adaa8764138aac2e44f3af1", size = 23784, upload-time = "2025-09-27T18:36:47.884Z" }, + { url = "https://files.pythonhosted.org/packages/ed/76/104b2aa106a208da8b17a2fb72e033a5a9d7073c68f7e508b94916ed47a9/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:795e7751525cae078558e679d646ae45574b47ed6e7771863fcc079a6171a0fc", size = 21588, upload-time = "2025-09-27T18:36:48.82Z" }, + { url = "https://files.pythonhosted.org/packages/b5/99/16a5eb2d140087ebd97180d95249b00a03aa87e29cc224056274f2e45fd6/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8485f406a96febb5140bfeca44a73e3ce5116b2501ac54fe953e488fb1d03b12", size = 23041, upload-time = "2025-09-27T18:36:49.797Z" }, + { url = "https://files.pythonhosted.org/packages/19/bc/e7140ed90c5d61d77cea142eed9f9c303f4c4806f60a1044c13e3f1471d0/markupsafe-3.0.3-cp313-cp313-win32.whl", hash = "sha256:bdd37121970bfd8be76c5fb069c7751683bdf373db1ed6c010162b2a130248ed", size = 14543, upload-time = "2025-09-27T18:36:51.584Z" }, + { url = "https://files.pythonhosted.org/packages/05/73/c4abe620b841b6b791f2edc248f556900667a5a1cf023a6646967ae98335/markupsafe-3.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:9a1abfdc021a164803f4d485104931fb8f8c1efd55bc6b748d2f5774e78b62c5", size = 15113, upload-time = "2025-09-27T18:36:52.537Z" }, + { url = "https://files.pythonhosted.org/packages/f0/3a/fa34a0f7cfef23cf9500d68cb7c32dd64ffd58a12b09225fb03dd37d5b80/markupsafe-3.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:7e68f88e5b8799aa49c85cd116c932a1ac15caaa3f5db09087854d218359e485", size = 13911, upload-time = "2025-09-27T18:36:53.513Z" }, + { url = "https://files.pythonhosted.org/packages/e4/d7/e05cd7efe43a88a17a37b3ae96e79a19e846f3f456fe79c57ca61356ef01/markupsafe-3.0.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:218551f6df4868a8d527e3062d0fb968682fe92054e89978594c28e642c43a73", size = 11658, upload-time = "2025-09-27T18:36:54.819Z" }, + { url = "https://files.pythonhosted.org/packages/99/9e/e412117548182ce2148bdeacdda3bb494260c0b0184360fe0d56389b523b/markupsafe-3.0.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:3524b778fe5cfb3452a09d31e7b5adefeea8c5be1d43c4f810ba09f2ceb29d37", size = 12066, upload-time = "2025-09-27T18:36:55.714Z" }, + { url = "https://files.pythonhosted.org/packages/bc/e6/fa0ffcda717ef64a5108eaa7b4f5ed28d56122c9a6d70ab8b72f9f715c80/markupsafe-3.0.3-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4e885a3d1efa2eadc93c894a21770e4bc67899e3543680313b09f139e149ab19", size = 25639, upload-time = "2025-09-27T18:36:56.908Z" }, + { url = "https://files.pythonhosted.org/packages/96/ec/2102e881fe9d25fc16cb4b25d5f5cde50970967ffa5dddafdb771237062d/markupsafe-3.0.3-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8709b08f4a89aa7586de0aadc8da56180242ee0ada3999749b183aa23df95025", size = 23569, upload-time = "2025-09-27T18:36:57.913Z" }, + { url = "https://files.pythonhosted.org/packages/4b/30/6f2fce1f1f205fc9323255b216ca8a235b15860c34b6798f810f05828e32/markupsafe-3.0.3-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b8512a91625c9b3da6f127803b166b629725e68af71f8184ae7e7d54686a56d6", size = 23284, upload-time = "2025-09-27T18:36:58.833Z" }, + { url = "https://files.pythonhosted.org/packages/58/47/4a0ccea4ab9f5dcb6f79c0236d954acb382202721e704223a8aafa38b5c8/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:9b79b7a16f7fedff2495d684f2b59b0457c3b493778c9eed31111be64d58279f", size = 24801, upload-time = "2025-09-27T18:36:59.739Z" }, + { url = "https://files.pythonhosted.org/packages/6a/70/3780e9b72180b6fecb83a4814d84c3bf4b4ae4bf0b19c27196104149734c/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:12c63dfb4a98206f045aa9563db46507995f7ef6d83b2f68eda65c307c6829eb", size = 22769, upload-time = "2025-09-27T18:37:00.719Z" }, + { url = "https://files.pythonhosted.org/packages/98/c5/c03c7f4125180fc215220c035beac6b9cb684bc7a067c84fc69414d315f5/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8f71bc33915be5186016f675cd83a1e08523649b0e33efdb898db577ef5bb009", size = 23642, upload-time = "2025-09-27T18:37:01.673Z" }, + { url = "https://files.pythonhosted.org/packages/80/d6/2d1b89f6ca4bff1036499b1e29a1d02d282259f3681540e16563f27ebc23/markupsafe-3.0.3-cp313-cp313t-win32.whl", hash = "sha256:69c0b73548bc525c8cb9a251cddf1931d1db4d2258e9599c28c07ef3580ef354", size = 14612, upload-time = "2025-09-27T18:37:02.639Z" }, + { url = "https://files.pythonhosted.org/packages/2b/98/e48a4bfba0a0ffcf9925fe2d69240bfaa19c6f7507b8cd09c70684a53c1e/markupsafe-3.0.3-cp313-cp313t-win_amd64.whl", hash = "sha256:1b4b79e8ebf6b55351f0d91fe80f893b4743f104bff22e90697db1590e47a218", size = 15200, upload-time = "2025-09-27T18:37:03.582Z" }, + { url = "https://files.pythonhosted.org/packages/0e/72/e3cc540f351f316e9ed0f092757459afbc595824ca724cbc5a5d4263713f/markupsafe-3.0.3-cp313-cp313t-win_arm64.whl", hash = "sha256:ad2cf8aa28b8c020ab2fc8287b0f823d0a7d8630784c31e9ee5edea20f406287", size = 13973, upload-time = "2025-09-27T18:37:04.929Z" }, + { url = "https://files.pythonhosted.org/packages/33/8a/8e42d4838cd89b7dde187011e97fe6c3af66d8c044997d2183fbd6d31352/markupsafe-3.0.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:eaa9599de571d72e2daf60164784109f19978b327a3910d3e9de8c97b5b70cfe", size = 11619, upload-time = "2025-09-27T18:37:06.342Z" }, + { url = "https://files.pythonhosted.org/packages/b5/64/7660f8a4a8e53c924d0fa05dc3a55c9cee10bbd82b11c5afb27d44b096ce/markupsafe-3.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c47a551199eb8eb2121d4f0f15ae0f923d31350ab9280078d1e5f12b249e0026", size = 12029, upload-time = "2025-09-27T18:37:07.213Z" }, + { url = "https://files.pythonhosted.org/packages/da/ef/e648bfd021127bef5fa12e1720ffed0c6cbb8310c8d9bea7266337ff06de/markupsafe-3.0.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f34c41761022dd093b4b6896d4810782ffbabe30f2d443ff5f083e0cbbb8c737", size = 24408, upload-time = "2025-09-27T18:37:09.572Z" }, + { url = "https://files.pythonhosted.org/packages/41/3c/a36c2450754618e62008bf7435ccb0f88053e07592e6028a34776213d877/markupsafe-3.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:457a69a9577064c05a97c41f4e65148652db078a3a509039e64d3467b9e7ef97", size = 23005, upload-time = "2025-09-27T18:37:10.58Z" }, + { url = "https://files.pythonhosted.org/packages/bc/20/b7fdf89a8456b099837cd1dc21974632a02a999ec9bf7ca3e490aacd98e7/markupsafe-3.0.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e8afc3f2ccfa24215f8cb28dcf43f0113ac3c37c2f0f0806d8c70e4228c5cf4d", size = 22048, upload-time = "2025-09-27T18:37:11.547Z" }, + { url = "https://files.pythonhosted.org/packages/9a/a7/591f592afdc734f47db08a75793a55d7fbcc6902a723ae4cfbab61010cc5/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ec15a59cf5af7be74194f7ab02d0f59a62bdcf1a537677ce67a2537c9b87fcda", size = 23821, upload-time = "2025-09-27T18:37:12.48Z" }, + { url = "https://files.pythonhosted.org/packages/7d/33/45b24e4f44195b26521bc6f1a82197118f74df348556594bd2262bda1038/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:0eb9ff8191e8498cca014656ae6b8d61f39da5f95b488805da4bb029cccbfbaf", size = 21606, upload-time = "2025-09-27T18:37:13.485Z" }, + { url = "https://files.pythonhosted.org/packages/ff/0e/53dfaca23a69fbfbbf17a4b64072090e70717344c52eaaaa9c5ddff1e5f0/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2713baf880df847f2bece4230d4d094280f4e67b1e813eec43b4c0e144a34ffe", size = 23043, upload-time = "2025-09-27T18:37:14.408Z" }, + { url = "https://files.pythonhosted.org/packages/46/11/f333a06fc16236d5238bfe74daccbca41459dcd8d1fa952e8fbd5dccfb70/markupsafe-3.0.3-cp314-cp314-win32.whl", hash = "sha256:729586769a26dbceff69f7a7dbbf59ab6572b99d94576a5592625d5b411576b9", size = 14747, upload-time = "2025-09-27T18:37:15.36Z" }, + { url = "https://files.pythonhosted.org/packages/28/52/182836104b33b444e400b14f797212f720cbc9ed6ba34c800639d154e821/markupsafe-3.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:bdc919ead48f234740ad807933cdf545180bfbe9342c2bb451556db2ed958581", size = 15341, upload-time = "2025-09-27T18:37:16.496Z" }, + { url = "https://files.pythonhosted.org/packages/6f/18/acf23e91bd94fd7b3031558b1f013adfa21a8e407a3fdb32745538730382/markupsafe-3.0.3-cp314-cp314-win_arm64.whl", hash = "sha256:5a7d5dc5140555cf21a6fefbdbf8723f06fcd2f63ef108f2854de715e4422cb4", size = 14073, upload-time = "2025-09-27T18:37:17.476Z" }, + { url = "https://files.pythonhosted.org/packages/3c/f0/57689aa4076e1b43b15fdfa646b04653969d50cf30c32a102762be2485da/markupsafe-3.0.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:1353ef0c1b138e1907ae78e2f6c63ff67501122006b0f9abad68fda5f4ffc6ab", size = 11661, upload-time = "2025-09-27T18:37:18.453Z" }, + { url = "https://files.pythonhosted.org/packages/89/c3/2e67a7ca217c6912985ec766c6393b636fb0c2344443ff9d91404dc4c79f/markupsafe-3.0.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1085e7fbddd3be5f89cc898938f42c0b3c711fdcb37d75221de2666af647c175", size = 12069, upload-time = "2025-09-27T18:37:19.332Z" }, + { url = "https://files.pythonhosted.org/packages/f0/00/be561dce4e6ca66b15276e184ce4b8aec61fe83662cce2f7d72bd3249d28/markupsafe-3.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1b52b4fb9df4eb9ae465f8d0c228a00624de2334f216f178a995ccdcf82c4634", size = 25670, upload-time = "2025-09-27T18:37:20.245Z" }, + { url = "https://files.pythonhosted.org/packages/50/09/c419f6f5a92e5fadde27efd190eca90f05e1261b10dbd8cbcb39cd8ea1dc/markupsafe-3.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fed51ac40f757d41b7c48425901843666a6677e3e8eb0abcff09e4ba6e664f50", size = 23598, upload-time = "2025-09-27T18:37:21.177Z" }, + { url = "https://files.pythonhosted.org/packages/22/44/a0681611106e0b2921b3033fc19bc53323e0b50bc70cffdd19f7d679bb66/markupsafe-3.0.3-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f190daf01f13c72eac4efd5c430a8de82489d9cff23c364c3ea822545032993e", size = 23261, upload-time = "2025-09-27T18:37:22.167Z" }, + { url = "https://files.pythonhosted.org/packages/5f/57/1b0b3f100259dc9fffe780cfb60d4be71375510e435efec3d116b6436d43/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e56b7d45a839a697b5eb268c82a71bd8c7f6c94d6fd50c3d577fa39a9f1409f5", size = 24835, upload-time = "2025-09-27T18:37:23.296Z" }, + { url = "https://files.pythonhosted.org/packages/26/6a/4bf6d0c97c4920f1597cc14dd720705eca0bf7c787aebc6bb4d1bead5388/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:f3e98bb3798ead92273dc0e5fd0f31ade220f59a266ffd8a4f6065e0a3ce0523", size = 22733, upload-time = "2025-09-27T18:37:24.237Z" }, + { url = "https://files.pythonhosted.org/packages/14/c7/ca723101509b518797fedc2fdf79ba57f886b4aca8a7d31857ba3ee8281f/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5678211cb9333a6468fb8d8be0305520aa073f50d17f089b5b4b477ea6e67fdc", size = 23672, upload-time = "2025-09-27T18:37:25.271Z" }, + { url = "https://files.pythonhosted.org/packages/fb/df/5bd7a48c256faecd1d36edc13133e51397e41b73bb77e1a69deab746ebac/markupsafe-3.0.3-cp314-cp314t-win32.whl", hash = "sha256:915c04ba3851909ce68ccc2b8e2cd691618c4dc4c4232fb7982bca3f41fd8c3d", size = 14819, upload-time = "2025-09-27T18:37:26.285Z" }, + { url = "https://files.pythonhosted.org/packages/1a/8a/0402ba61a2f16038b48b39bccca271134be00c5c9f0f623208399333c448/markupsafe-3.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:4faffd047e07c38848ce017e8725090413cd80cbc23d86e55c587bf979e579c9", size = 15426, upload-time = "2025-09-27T18:37:27.316Z" }, + { url = "https://files.pythonhosted.org/packages/70/bc/6f1c2f612465f5fa89b95bead1f44dcb607670fd42891d8fdcd5d039f4f4/markupsafe-3.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:32001d6a8fc98c8cb5c947787c5d08b0a50663d139f1305bac5885d98d9b40fa", size = 14146, upload-time = "2025-09-27T18:37:28.327Z" }, +] + +[[package]] +name = "mpmath" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, +] + +[[package]] +name = "nanobind" +version = "2.15.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0d/08/cef3343279d2bde0baa57aadaacf7196ee336f12449c50c27f799c45b3f7/nanobind-2.15.0.tar.gz", hash = "sha256:3eace37a3b8cdadd531fcc6145263985b29821467161d34ff722b0644cf445b5", size = 1060056, upload-time = "2026-08-15T19:01:04.369Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d8/72/16a606cbea58be568475f450ece121b3fd62cffb390fd0f8eb8b410904f5/nanobind-2.15.0-py3-none-any.whl", hash = "sha256:91dae9305c6242cf63ab11a5c41594e234e7a22fd9fcc2c5c4ba80724abc1787", size = 271648, upload-time = "2026-08-15T19:01:02.477Z" }, +] + +[[package]] +name = "networkx" +version = "3.6.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" }, +] + +[[package]] +name = "ninja" +version = "1.13.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/43/73/79a0b22fc731989c708068427579e840a6cf4e937fe7ae5c5d0b7356ac22/ninja-1.13.0.tar.gz", hash = "sha256:4a40ce995ded54d9dc24f8ea37ff3bf62ad192b547f6c7126e7e25045e76f978", size = 242558, upload-time = "2025-08-11T15:10:19.421Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3c/74/d02409ed2aa865e051b7edda22ad416a39d81a84980f544f8de717cab133/ninja-1.13.0-py3-none-macosx_10_9_universal2.whl", hash = "sha256:fa2a8bfc62e31b08f83127d1613d10821775a0eb334197154c4d6067b7068ff1", size = 310125, upload-time = "2025-08-11T15:09:50.971Z" }, + { url = "https://files.pythonhosted.org/packages/8e/de/6e1cd6b84b412ac1ef327b76f0641aeb5dcc01e9d3f9eee0286d0c34fd93/ninja-1.13.0-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3d00c692fb717fd511abeb44b8c5d00340c36938c12d6538ba989fe764e79630", size = 177467, upload-time = "2025-08-11T15:09:52.767Z" }, + { url = "https://files.pythonhosted.org/packages/c8/83/49320fb6e58ae3c079381e333575fdbcf1cca3506ee160a2dcce775046fa/ninja-1.13.0-py3-none-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:be7f478ff9f96a128b599a964fc60a6a87b9fa332ee1bd44fa243ac88d50291c", size = 187834, upload-time = "2025-08-11T15:09:54.115Z" }, + { url = "https://files.pythonhosted.org/packages/56/c7/ba22748fb59f7f896b609cd3e568d28a0a367a6d953c24c461fe04fc4433/ninja-1.13.0-py3-none-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:60056592cf495e9a6a4bea3cd178903056ecb0943e4de45a2ea825edb6dc8d3e", size = 202736, upload-time = "2025-08-11T15:09:55.745Z" }, + { url = "https://files.pythonhosted.org/packages/79/22/d1de07632b78ac8e6b785f41fa9aad7a978ec8c0a1bf15772def36d77aac/ninja-1.13.0-py3-none-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:1c97223cdda0417f414bf864cfb73b72d8777e57ebb279c5f6de368de0062988", size = 179034, upload-time = "2025-08-11T15:09:57.394Z" }, + { url = "https://files.pythonhosted.org/packages/ed/de/0e6edf44d6a04dabd0318a519125ed0415ce437ad5a1ec9b9be03d9048cf/ninja-1.13.0-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fb46acf6b93b8dd0322adc3a4945452a4e774b75b91293bafcc7b7f8e6517dfa", size = 180716, upload-time = "2025-08-11T15:09:58.696Z" }, + { url = "https://files.pythonhosted.org/packages/54/28/938b562f9057aaa4d6bfbeaa05e81899a47aebb3ba6751e36c027a7f5ff7/ninja-1.13.0-py3-none-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:4be9c1b082d244b1ad7ef41eb8ab088aae8c109a9f3f0b3e56a252d3e00f42c1", size = 146843, upload-time = "2025-08-11T15:10:00.046Z" }, + { url = "https://files.pythonhosted.org/packages/2a/fb/d06a3838de4f8ab866e44ee52a797b5491df823901c54943b2adb0389fbb/ninja-1.13.0-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:6739d3352073341ad284246f81339a384eec091d9851a886dfa5b00a6d48b3e2", size = 154402, upload-time = "2025-08-11T15:10:01.657Z" }, + { url = "https://files.pythonhosted.org/packages/31/bf/0d7808af695ceddc763cf251b84a9892cd7f51622dc8b4c89d5012779f06/ninja-1.13.0-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:11be2d22027bde06f14c343f01d31446747dbb51e72d00decca2eb99be911e2f", size = 552388, upload-time = "2025-08-11T15:10:03.349Z" }, + { url = "https://files.pythonhosted.org/packages/9d/70/c99d0c2c809f992752453cce312848abb3b1607e56d4cd1b6cded317351a/ninja-1.13.0-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:aa45b4037b313c2f698bc13306239b8b93b4680eb47e287773156ac9e9304714", size = 472501, upload-time = "2025-08-11T15:10:04.735Z" }, + { url = "https://files.pythonhosted.org/packages/9f/43/c217b1153f0e499652f5e0766da8523ce3480f0a951039c7af115e224d55/ninja-1.13.0-py3-none-musllinux_1_2_i686.whl", hash = "sha256:5f8e1e8a1a30835eeb51db05cf5a67151ad37542f5a4af2a438e9490915e5b72", size = 638280, upload-time = "2025-08-11T15:10:06.512Z" }, + { url = "https://files.pythonhosted.org/packages/8c/45/9151bba2c8d0ae2b6260f71696330590de5850e5574b7b5694dce6023e20/ninja-1.13.0-py3-none-musllinux_1_2_ppc64le.whl", hash = "sha256:3d7d7779d12cb20c6d054c61b702139fd23a7a964ec8f2c823f1ab1b084150db", size = 642420, upload-time = "2025-08-11T15:10:08.35Z" }, + { url = "https://files.pythonhosted.org/packages/3c/fb/95752eb635bb8ad27d101d71bef15bc63049de23f299e312878fc21cb2da/ninja-1.13.0-py3-none-musllinux_1_2_riscv64.whl", hash = "sha256:d741a5e6754e0bda767e3274a0f0deeef4807f1fec6c0d7921a0244018926ae5", size = 585106, upload-time = "2025-08-11T15:10:09.818Z" }, + { url = "https://files.pythonhosted.org/packages/c1/31/aa56a1a286703800c0cbe39fb4e82811c277772dc8cd084f442dd8e2938a/ninja-1.13.0-py3-none-musllinux_1_2_s390x.whl", hash = "sha256:e8bad11f8a00b64137e9b315b137d8bb6cbf3086fbdc43bf1f90fd33324d2e96", size = 707138, upload-time = "2025-08-11T15:10:11.366Z" }, + { url = "https://files.pythonhosted.org/packages/34/6f/5f5a54a1041af945130abdb2b8529cbef0cdcbbf9bcf3f4195378319d29a/ninja-1.13.0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:b4f2a072db3c0f944c32793e91532d8948d20d9ab83da9c0c7c15b5768072200", size = 581758, upload-time = "2025-08-11T15:10:13.295Z" }, + { url = "https://files.pythonhosted.org/packages/95/97/51359c77527d45943fe7a94d00a3843b81162e6c4244b3579fe8fc54cb9c/ninja-1.13.0-py3-none-win32.whl", hash = "sha256:8cfbb80b4a53456ae8a39f90ae3d7a2129f45ea164f43fadfa15dc38c4aef1c9", size = 267201, upload-time = "2025-08-11T15:10:15.158Z" }, + { url = "https://files.pythonhosted.org/packages/29/45/c0adfbfb0b5895aa18cec400c535b4f7ff3e52536e0403602fc1a23f7de9/ninja-1.13.0-py3-none-win_amd64.whl", hash = "sha256:fb8ee8719f8af47fed145cced4a85f0755dd55d45b2bddaf7431fa89803c5f3e", size = 309975, upload-time = "2025-08-11T15:10:16.697Z" }, + { url = "https://files.pythonhosted.org/packages/df/93/a7b983643d1253bb223234b5b226e69de6cda02b76cdca7770f684b795f5/ninja-1.13.0-py3-none-win_arm64.whl", hash = "sha256:3c0b40b1f0bba764644385319028650087b4c1b18cdfa6f45cb39a3669b81aa9", size = 290806, upload-time = "2025-08-11T15:10:18.018Z" }, +] + +[[package]] +name = "numpy" +version = "2.4.6" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.12' and sys_platform != 'darwin'", + "python_full_version < '3.12' and sys_platform == 'darwin'", +] +sdist = { url = "https://files.pythonhosted.org/packages/d0/ad/fed0499ce6a338d2a03ebae59cd15093910c8875328855781952abf6c2fe/numpy-2.4.6.tar.gz", hash = "sha256:f3a3570c4a2a16746ac2c31a7c7c7b0c186b95ce902e33db6f28094ed7387dda", size = 20735807, upload-time = "2026-05-18T23:37:14.07Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/49/ec46835a70be8fa6446c495126ac84fdb28cb2558e1620ffb87a10c8b64c/numpy-2.4.6-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:0280e0356c0829a18d9de1cb7eee50ec22ca639878d7240307ca0943d73cd2c4", size = 16969194, upload-time = "2026-05-18T23:33:13.503Z" }, + { url = "https://files.pythonhosted.org/packages/0e/0d/f5957185c0ee2f3e12f78715aa9e3b353fd83633316c8532b38faa37e3f6/numpy-2.4.6-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:110f8b71aacb688ec69062bb7f6938a0f8acb01b7c1c4beb453c65b6d234584d", size = 14964111, upload-time = "2026-05-18T23:33:17.795Z" }, + { url = "https://files.pythonhosted.org/packages/ad/40/40a40ee0ddf7ceb782c49af278894b686e586d65d8c1889c8b5da01a3d7d/numpy-2.4.6-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:4cfe66903cc32a9921a6733d96b19bb6abf310397581bbad89c228f5abaf0ee8", size = 5469159, upload-time = "2026-05-18T23:33:20.654Z" }, + { url = "https://files.pythonhosted.org/packages/63/13/f9a8046535cb21deae82f8d03de9617e08882d274fad2539630761888228/numpy-2.4.6-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:8155154c7c691289fe18f510b5d4657c68c67989f293f0535a91360392ff6538", size = 6798936, upload-time = "2026-05-18T23:33:22.987Z" }, + { url = "https://files.pythonhosted.org/packages/33/a8/6fa8c1a345a8c85dbb21932c447bee07c30a2c2a3f31e369c0a84b300147/numpy-2.4.6-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0ab0a9c4ffb1a6d95ef519fe4247dba8eb6b18ad93999f76b7f657039acabd47", size = 15966692, upload-time = "2026-05-18T23:33:26.62Z" }, + { url = "https://files.pythonhosted.org/packages/02/03/74fe2a4cb3817d94d86402f2506554130a2f01414e299b5a843e5a8a957f/numpy-2.4.6-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:89cd468399cfd2504718f0ba50e410dca55a170b61a02ad92bb18c8a65186e93", size = 16918164, upload-time = "2026-05-18T23:33:29.955Z" }, + { url = "https://files.pythonhosted.org/packages/c5/80/3615be3313f7e7696609bc194b9f0101da809df79e859bdb84e0cd043f46/numpy-2.4.6-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:c2d37ab77531417474168eb79d6d80b14f821a966818505d03013d0833edb7a8", size = 17322877, upload-time = "2026-05-18T23:33:34.724Z" }, + { url = "https://files.pythonhosted.org/packages/ca/ac/a691e0fe2675e370d0e08ff905adc49a1c8830e8cae03efe4477e92cd55d/numpy-2.4.6-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:f407cb6b8e9d6d8c626bc73c945db1706035af8fd632295547bf1c9e46d092d6", size = 18651487, upload-time = "2026-05-18T23:33:38.217Z" }, + { url = "https://files.pythonhosted.org/packages/15/a7/9bc1cd626d7bf6869bfedf27b91b6ab5dd607758bf8e959d6fa80c6a59cb/numpy-2.4.6-cp311-cp311-win32.whl", hash = "sha256:ddea102b48f9e339f3948bf22040944184627a30fdf7f858667673b9c5f033c8", size = 6233945, upload-time = "2026-05-18T23:33:41.331Z" }, + { url = "https://files.pythonhosted.org/packages/c5/31/7fc6239c12bce7e931463251cca4426c465e1876ba3cc785402ef4dd8f4e/numpy-2.4.6-cp311-cp311-win_amd64.whl", hash = "sha256:1e254a00cdf42b1e4d5b3d68d33af63268d41340d8885df2ab6470f2e1500147", size = 12608406, upload-time = "2026-05-18T23:33:44.131Z" }, + { url = "https://files.pythonhosted.org/packages/27/83/140f85a466595a16382996a1bf06b2b54bcd597488921b0c9daaeeda72af/numpy-2.4.6-cp311-cp311-win_arm64.whl", hash = "sha256:ed9749eef4cbd126da3dc1d6bcb3a57f5eb7ac6a6484146bdbf743f552dfc577", size = 10479528, upload-time = "2026-05-18T23:33:50.725Z" }, + { url = "https://files.pythonhosted.org/packages/95/2a/3d7b5ac8aac24feaf9ad7ed58f45b0bbc06d37e4338ae84c9f2298b570f9/numpy-2.4.6-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:001fbb8e08d942dd57599e781f2472269ee7f2755fae407b4f67b2f0b17da3f1", size = 16689119, upload-time = "2026-05-18T23:33:54.065Z" }, + { url = "https://files.pythonhosted.org/packages/ea/12/92c4c131527599e8288d6918e888d88726f84d805d784b771f32408aeaef/numpy-2.4.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ebfb099f8dcf083deef3ac1ca4c1503f387cf76296fcb3816b66f5ecb5f54fdb", size = 14699246, upload-time = "2026-05-18T23:33:57.621Z" }, + { url = "https://files.pythonhosted.org/packages/ad/fe/c0a6b7b2ca128a8fb228575147073b660656734b8ebe4d76c8fd748dcc79/numpy-2.4.6-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:3213d622a0283a39a93d188f3cf72b26862df52fbb4ca3697f51705016523d41", size = 5204410, upload-time = "2026-05-18T23:34:00.302Z" }, + { url = "https://files.pythonhosted.org/packages/f3/d4/9770d14ba719432bb90a421bfd443872ed0f70f7264b64bec12ea363d5fd/numpy-2.4.6-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:357cc07a6d7b0b182ff02249616a03742827ebb1277546b5c7cd7f7620a45698", size = 6551240, upload-time = "2026-05-18T23:34:02.852Z" }, + { url = "https://files.pythonhosted.org/packages/c9/c6/50a46a6205feba2343f1d6d17438107c5dc491ed1c736e6ea68689fd906b/numpy-2.4.6-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5f9fb9157b4ce2971008323afe46053787b526ef624fea915b261468a8421a0f", size = 15671012, upload-time = "2026-05-18T23:34:05.485Z" }, + { url = "https://files.pythonhosted.org/packages/99/60/14115e6364fa676c5397c2ad3004e527e9aa487abf5d0706ec81bbd08529/numpy-2.4.6-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:90f9849678c75fe7afa2d348ac842c168b0a4d3d61919687216dfc547976d853", size = 16645538, upload-time = "2026-05-18T23:34:09.265Z" }, + { url = "https://files.pythonhosted.org/packages/ae/c5/693cbe59e57db94d2231fa519ca3978dc9e19da5a8f088588f5c6e947ff2/numpy-2.4.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:c1a2af6c6ef86344a6b0db6b97834208bf598db514f2b155042439b62605601a", size = 17020706, upload-time = "2026-05-18T23:34:13.053Z" }, + { url = "https://files.pythonhosted.org/packages/ef/fc/85b7c4eff9b4966ade25c2273cf7e7012e92366c032058653934b37de044/numpy-2.4.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:e5805d5a22fd19c8ccff10a9561f9df94436b0545619ea579db2d3c35294bce2", size = 18368541, upload-time = "2026-05-18T23:34:17.024Z" }, + { url = "https://files.pythonhosted.org/packages/f6/81/e1b27545deedce7f4a0b348618c6b62d74e36a4dc9ccd42f3eb2f85eee32/numpy-2.4.6-cp312-cp312-win32.whl", hash = "sha256:e3eeb0aabd6bd5ce64faae67e9935203a6991b4bc2a485a767fbafb2c5125f45", size = 5962825, upload-time = "2026-05-18T23:34:20.3Z" }, + { url = "https://files.pythonhosted.org/packages/ab/ca/feab00bd44aa5fe1ad2c18f08b4d3bb92e26484b0b1d1443897809ed528c/numpy-2.4.6-cp312-cp312-win_amd64.whl", hash = "sha256:d8e8286dd7cea7895157318d1b91cdacac64c479f3cbc8dce548331728484751", size = 12321687, upload-time = "2026-05-18T23:34:23.095Z" }, + { url = "https://files.pythonhosted.org/packages/63/cf/5a6d34850a39d1093558564f77ee8e8e0bee5061151b8f05a55711001ec7/numpy-2.4.6-cp312-cp312-win_arm64.whl", hash = "sha256:4081eb135ac24158bd51cdfbef16f1c64df7063b1143f24731387137c092bec8", size = 10221482, upload-time = "2026-05-18T23:34:25.876Z" }, + { url = "https://files.pythonhosted.org/packages/fb/82/bdab26d7438c6791ca31b7c024ca37c1eab8b726ba236129005cd4a06e45/numpy-2.4.6-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:511dbaf848decaaaf4b4ca48032619fb3138710c4bf7da7617765edad1ef96b0", size = 16684648, upload-time = "2026-05-18T23:34:29.41Z" }, + { url = "https://files.pythonhosted.org/packages/1b/30/a80189bcc7f5e4258b3fbc3968d909d1756f54d023299ecc39ad6fdb9ef8/numpy-2.4.6-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:bf162abab1c1a736333192707cef898e735a5ca00f38f27eeedf44b39d9e85eb", size = 14693902, upload-time = "2026-05-18T23:34:33.013Z" }, + { url = "https://files.pythonhosted.org/packages/97/12/70b5d0d7c15e1ebb8a6a84a8caa1d19e181d84fb58bb6d70aca29099dec1/numpy-2.4.6-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:043191bfa8eab18c776647b62723ac9dddece59743b13f49b2016094129c2b3f", size = 5198992, upload-time = "2026-05-18T23:34:36.132Z" }, + { url = "https://files.pythonhosted.org/packages/ba/8c/ebd2a8f8a83541f8d38cc5667e8c2b69cecfd30da6e45693e8158857d44b/numpy-2.4.6-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:6180d8b35af935aed8ece3a85e0a43f87393ae0ac87c8d2c8bd2c993f7270ef3", size = 6546944, upload-time = "2026-05-18T23:34:38.484Z" }, + { url = "https://files.pythonhosted.org/packages/bb/c5/7b863a97a91671a0338f4253bd3b5a3d3852f0692dae91711c9f4a10e787/numpy-2.4.6-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:72fbe16c6fac95aedf5937fa873445cec2110be35d8a4e9433d7501fd98dae6b", size = 15669392, upload-time = "2026-05-18T23:34:41.257Z" }, + { url = "https://files.pythonhosted.org/packages/a5/9d/3584b9984ca4c047aea75214ce1a4c4c73d849bd71b604264b7f5653f8a8/numpy-2.4.6-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a7830bab239b79cda9c08c2da014761cafb48da6150e1da17ac06283f43b6089", size = 16633220, upload-time = "2026-05-18T23:34:45.075Z" }, + { url = "https://files.pythonhosted.org/packages/05/ae/7c67fba23bd98caec7c99261f3a16072ade14813486b0282cb29846de832/numpy-2.4.6-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:ef4aea96ce4d3b074422cb4f2f64e216bf9e213004bb58ecfdf50ea02ea8eb9a", size = 17020800, upload-time = "2026-05-18T23:34:49.065Z" }, + { url = "https://files.pythonhosted.org/packages/d9/5d/3b6725cb31d983c5e66916f5d36f6d7e5521129e4c4404d64f918292a5b6/numpy-2.4.6-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:dfa20cc6ca228e6b155b11da03825975ce66aea520985dbbddf0f2a5a495c605", size = 18357600, upload-time = "2026-05-18T23:34:52.709Z" }, + { url = "https://files.pythonhosted.org/packages/f7/da/2ccc6c2fe8898dee01d90c75c5f5f914a23daf99e3e0f59516a08760c8b5/numpy-2.4.6-cp313-cp313-win32.whl", hash = "sha256:56b39e5e0622a09a25bf5baf62f4bcf0cb8a41ae6e2819cf49bbc5a74c083f91", size = 5961134, upload-time = "2026-05-18T23:34:55.618Z" }, + { url = "https://files.pythonhosted.org/packages/b5/cd/9cc4dc876fb065d5c220aae4d5e14826b2715331bb7618ce1fb07a679d99/numpy-2.4.6-cp313-cp313-win_amd64.whl", hash = "sha256:c4fc99836233ea196540b17ab0983aff60ed07941751930f5f4d05bc3b3b7359", size = 12318598, upload-time = "2026-05-18T23:34:58.928Z" }, + { url = "https://files.pythonhosted.org/packages/39/1e/c0bcba1f8694116485fe28fd1be698c278fcda4141c5b0e53a2aed8b12a8/numpy-2.4.6-cp313-cp313-win_arm64.whl", hash = "sha256:a7c711e21628b52034bb5ab8d1bce291f752fcc5e92accc615778acee1ff4778", size = 10222272, upload-time = "2026-05-18T23:35:02.167Z" }, + { url = "https://files.pythonhosted.org/packages/63/6d/cc5619247c8f4204e507f5883528372e4ac4bb189e579fb859a12e480b1f/numpy-2.4.6-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:112b06a867b235ef466ed3508ddf0238050df9c727cafb5301ac385b899189a1", size = 14821197, upload-time = "2026-05-18T23:35:05.468Z" }, + { url = "https://files.pythonhosted.org/packages/00/58/f1c39161c87d9e9bed660f1ed4bafc0e403d5ec9650b6dd77aead07d489b/numpy-2.4.6-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:eaf7fa2de5c0be8ae6ff8e9bea2ccd725e980541244521d8d4b5f3354a27babe", size = 5326287, upload-time = "2026-05-18T23:35:08.693Z" }, + { url = "https://files.pythonhosted.org/packages/af/57/3917ab0fd97f271a8694513581b8a36c655f111c446852c302f04ccdb6fc/numpy-2.4.6-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:7265a2f3d436e54ef9f2b52b5c937e6be778781bd97a590319d7348f1c1ca997", size = 6646763, upload-time = "2026-05-18T23:35:11.459Z" }, + { url = "https://files.pythonhosted.org/packages/eb/0f/037e64c494b67581ae18193d770adef354c41f3f2c8ebf865602d949bf8f/numpy-2.4.6-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f74a575920ab21fe304421a3fc28793d82e299cae9eccb37084e9fc7f3617c20", size = 15728070, upload-time = "2026-05-18T23:35:14.79Z" }, + { url = "https://files.pythonhosted.org/packages/21/a6/5d2bae9c9542eb4df16dc9c46dc79c186e9bad53805dfa5399a6023c6db0/numpy-2.4.6-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ede83e07a75dd06bc501566c1eca2afc0d61677c1472ac9ad93fdee6e638a48d", size = 16681752, upload-time = "2026-05-18T23:35:18.836Z" }, + { url = "https://files.pythonhosted.org/packages/92/14/23d1dfb410ae362cd59ce53e936b1513d545eb40db3949ced632e19a459e/numpy-2.4.6-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:68bb27509ac1b9a3443094260f6326150663b06abe40b73a2f81160623da5b67", size = 17086024, upload-time = "2026-05-18T23:35:22.52Z" }, + { url = "https://files.pythonhosted.org/packages/4b/6e/23595a2c642cdf3bc567877064bdd7f91c8b0038a4453cf2daf7248eafe9/numpy-2.4.6-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:a0df0043bdb289bde1f62da130d20df23d58b45429f752bc7a8fc5325a225ecd", size = 18403398, upload-time = "2026-05-18T23:35:26.398Z" }, + { url = "https://files.pythonhosted.org/packages/8a/90/0ac3bc947217e66dec77e7cbc6a1979d1af70b6461b82f620d3bccd5e4c8/numpy-2.4.6-cp313-cp313t-win32.whl", hash = "sha256:29a287e0cf63ff528da061de6b9f64a4618da591ca1046aafc54062e40ca7eab", size = 6084971, upload-time = "2026-05-18T23:35:29.387Z" }, + { url = "https://files.pythonhosted.org/packages/77/71/5673e351671a1d2bd6063b91b44f70c0affea7d1516fa7a6572941ba4aa1/numpy-2.4.6-cp313-cp313t-win_amd64.whl", hash = "sha256:25c692919ac5a01f170a3bfcd62d745b24fd095c353d50812637d6fcab442e75", size = 12458532, upload-time = "2026-05-18T23:35:32.175Z" }, + { url = "https://files.pythonhosted.org/packages/3f/88/19d3503c5046e688f049274b27a3ef3d771152fa80d3ba3d01a3dff61abe/numpy-2.4.6-cp313-cp313t-win_arm64.whl", hash = "sha256:1e978ec1e8bd0e0e4de6bb75de9d30cbb74db6b6a2bb727618613703ca0167dd", size = 10291881, upload-time = "2026-05-18T23:35:35.465Z" }, + { url = "https://files.pythonhosted.org/packages/f8/91/3ab2044d05fd16d343c5ac2e69b127f1b2854040dd20b193257c78028bd3/numpy-2.4.6-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:06ca2f61ec4385a07a6977c55ba998a4466c123642b4a32694d3128fce18c079", size = 16683458, upload-time = "2026-05-18T23:35:38.353Z" }, + { url = "https://files.pythonhosted.org/packages/8e/62/764ce66fa4147ae6d73071a3abf804ffe606f174618697c571acdf26a7c9/numpy-2.4.6-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:38efbc8de75c7a0fc1ac190162d892787f3f47b57cc291231aafee36b80982b7", size = 14704559, upload-time = "2026-05-18T23:35:42.14Z" }, + { url = "https://files.pythonhosted.org/packages/60/61/23f27c172f022e04025b7dc2367f4d63c1a398120607ec896228649a6f48/numpy-2.4.6-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:d581b735e177fdcdce6fed8e7e8880a3fb6ee4e3653a3ac6af01c6f4c03effc5", size = 5209716, upload-time = "2026-05-18T23:35:45.377Z" }, + { url = "https://files.pythonhosted.org/packages/03/71/21cf70dc6ea3e3acb95fc53a265b2fc248b981f0194ceb5b475271b8809d/numpy-2.4.6-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:0a041d3d761dc3c35cc56ce0351506a02bcbc25f7b169f652435141a17db9096", size = 6543947, upload-time = "2026-05-18T23:35:47.926Z" }, + { url = "https://files.pythonhosted.org/packages/d5/91/64288395ee1799bd2e0b04a305dce9666da90c961e1f3fe982a05ee1c036/numpy-2.4.6-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:40fdc1ae7125e518ea98e53e69a4ebc27e1fd50510c47b7ea130cf21e5e1d42b", size = 15685197, upload-time = "2026-05-18T23:35:50.863Z" }, + { url = "https://files.pythonhosted.org/packages/f3/eb/ebffaa97dc55502df69584a8f0dcf07f69a3e0b3e2323670a2722db9aa39/numpy-2.4.6-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a2c306dea656c12c68f51f4cea133cbe78ca7435eb28c735eac1d3ebe73be6e8", size = 16638245, upload-time = "2026-05-18T23:35:54.752Z" }, + { url = "https://files.pythonhosted.org/packages/b8/0b/54f9da33128d7e350fab89c7455902eeae70349ee52bddb448dc4a576f45/numpy-2.4.6-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:33111801a01c12a8a1e3721f0a9232f8cfc8ae2c6b7098167e6f623c6073f402", size = 17036587, upload-time = "2026-05-18T23:35:58.355Z" }, + { url = "https://files.pythonhosted.org/packages/b6/f0/fdebc1052db1cc37c64beb22072d67cd6d1c71adca1299f53dec2b5e20d3/numpy-2.4.6-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:ae506e6902902557576a26ff33eda8695e7ecb3cb36c3b573a0765dee114ebdb", size = 18363226, upload-time = "2026-05-18T23:36:02.845Z" }, + { url = "https://files.pythonhosted.org/packages/aa/b4/298628d98c72b57e57f7165ae6a481a1deaf6f3c28262a6e4c739c275930/numpy-2.4.6-cp314-cp314-win32.whl", hash = "sha256:aaf159caa35993cb1f56fb9b8e4610d35758e7ca005412eb1daa856a78c9c4b1", size = 6010196, upload-time = "2026-05-18T23:36:05.92Z" }, + { url = "https://files.pythonhosted.org/packages/df/ac/46de6dda46478f7942f839e094970be2d4a861e005c4b3bf07c92e291a09/numpy-2.4.6-cp314-cp314-win_amd64.whl", hash = "sha256:b507f5c4c1d508876d1819b6bf9a49d365b96320b5d4993426b33a23ca4b8261", size = 12450334, upload-time = "2026-05-18T23:36:09.107Z" }, + { url = "https://files.pythonhosted.org/packages/78/92/b8b798ac784102c0da830d2257d59358e3d3d90d1e2b3f2575dad976c5cf/numpy-2.4.6-cp314-cp314-win_arm64.whl", hash = "sha256:6f41ae150c4e32db4f3310cdaf64b1593a03dbabe29eec77fc9b50fe64061df6", size = 10495678, upload-time = "2026-05-18T23:36:12.766Z" }, + { url = "https://files.pythonhosted.org/packages/30/34/ec28d1aa8115971537c01469ab2011ee96827930f0a124de1000cc2a7ed7/numpy-2.4.6-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:ece3d2cfe132e7d51f44a832b303895e6f2d499c5e74dfbdb06ee246147a304a", size = 14823672, upload-time = "2026-05-18T23:36:16.473Z" }, + { url = "https://files.pythonhosted.org/packages/16/bd/f6d1fede4e54e8042a7ff97bb495510f3c220f94bcd9e8b228e87c92cc0d/numpy-2.4.6-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:e3e5193ef5a3dc73bceee50f7fdc2c90dbb76c42df8d8fae3d1067a583df579e", size = 5328731, upload-time = "2026-05-18T23:36:19.767Z" }, + { url = "https://files.pythonhosted.org/packages/f4/f0/e105b9e2fd728a9910103884decd6951d9dd73896b914a98d9a231de02ee/numpy-2.4.6-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:17f9ade344e7d9b464a084d69bcf18fc691cb1db67c62ed80820bf4926d78f0e", size = 6649805, upload-time = "2026-05-18T23:36:22.266Z" }, + { url = "https://files.pythonhosted.org/packages/82/dd/1206a7ca6ab15e3f02069707ca96222e202af681bb73756da7527f3cb837/numpy-2.4.6-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9cd5ffd25db4e7ba6a375693b3fc0fc1791ec636c17db3720da19bde7180ec43", size = 15730496, upload-time = "2026-05-18T23:36:25.713Z" }, + { url = "https://files.pythonhosted.org/packages/51/e7/38d3ea825dcab85a591734decb2f6c67caa7c8367d374df1a1c3842f9b07/numpy-2.4.6-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7d92c3819208a60205a12a245c91ad70cb0a85336659b19b834205573ac8456e", size = 16679616, upload-time = "2026-05-18T23:36:29.652Z" }, + { url = "https://files.pythonhosted.org/packages/93/b7/caabfdf53edf663e0b4eb74d7d405d83baef09eb5e83bcd32d601d72b93e/numpy-2.4.6-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e85b752a1e912b70eaad4fafbd4d1238007ab221de2009b9a2f5ae7461239895", size = 17085145, upload-time = "2026-05-18T23:36:33.449Z" }, + { url = "https://files.pythonhosted.org/packages/f9/45/68d7c33a6bcf3e5aa3bdbd57a367e6f615286dfd6482f97e8ffeb734306e/numpy-2.4.6-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:29cb7f67d10b479ff07c17d33e39f78c07f71c40ef30d63c153d340e96cd3fb4", size = 18403813, upload-time = "2026-05-18T23:36:37.369Z" }, + { url = "https://files.pythonhosted.org/packages/9c/50/0753655aa844c99cd9e018aacf76f130f1bd81d881bb74bc0aef5d73a8ba/numpy-2.4.6-cp314-cp314t-win32.whl", hash = "sha256:260a5d70215b61ab4fadf5c7baacd64821842975eea312125ed3c39a6391b063", size = 6156982, upload-time = "2026-05-18T23:36:40.817Z" }, + { url = "https://files.pythonhosted.org/packages/b2/d4/7c67becf668f973cb490cec3e98dfd799d866f9c989a54d355672cfa0db6/numpy-2.4.6-cp314-cp314t-win_amd64.whl", hash = "sha256:81a1cca95ed5bb92aa8b10dd2cdc9a0d3853a50fad926c28b5d7e8ea54389627", size = 12638908, upload-time = "2026-05-18T23:36:43.996Z" }, + { url = "https://files.pythonhosted.org/packages/43/bb/e1c71a4295b1b1d1393d50dbb4f2a36283c6859d9d3892e84f00ec5a91d5/numpy-2.4.6-cp314-cp314t-win_arm64.whl", hash = "sha256:0c9136e14ed34a9e343a31c533d78a9813a69a3148332bce5e9821cb2f996e66", size = 10565867, upload-time = "2026-05-18T23:36:47.114Z" }, + { url = "https://files.pythonhosted.org/packages/de/12/b422cc84439adc0d00de605bf4a308890ae5c26f2c71fbd73e5d08fbb0dd/numpy-2.4.6-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:55cced7c52e981362f708ad635198e97a752dfba412cc03c23bbf3bd8d5cd662", size = 16847511, upload-time = "2026-05-18T23:36:50.673Z" }, + { url = "https://files.pythonhosted.org/packages/44/53/f481bef68011740f8849418d82db07230e825013f31f4eef5ba5b805316a/numpy-2.4.6-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:d6da64deb6b8ed903e7560180a92f2d804ee1ba5eeb849ac2748b8c1aba1f6d7", size = 14889064, upload-time = "2026-05-18T23:36:53.879Z" }, + { url = "https://files.pythonhosted.org/packages/7f/57/42ed575c10ced8af951d426bc4e1f8aff16fd851db33f067036215a7f860/numpy-2.4.6-pp311-pypy311_pp73-macosx_14_0_arm64.whl", hash = "sha256:68a5124b13fa6cc2086764a20005d30bc0548146f7f5322f02fce212ca14317f", size = 5394157, upload-time = "2026-05-18T23:36:57.194Z" }, + { url = "https://files.pythonhosted.org/packages/6a/ef/f66cc724fcc36c1e364c67f51ae9146090b8b584f27d58b97fdae3edd737/numpy-2.4.6-pp311-pypy311_pp73-macosx_14_0_x86_64.whl", hash = "sha256:948424b06129ce883307e8cff868c31396d8dc7630a59c61d70d98dbe70f222c", size = 6708728, upload-time = "2026-05-18T23:36:59.575Z" }, + { url = "https://files.pythonhosted.org/packages/1a/9c/c531f2293b91265d8b48e9b329f54fdd7ffae73cb4134ea10cca4237e9cc/numpy-2.4.6-pp311-pypy311_pp73-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5dbbdb29840ca3d91ee0fece42fc29278886d908280bfec0a5846c6f901a3eb0", size = 15798374, upload-time = "2026-05-18T23:37:02.674Z" }, + { url = "https://files.pythonhosted.org/packages/1a/b0/413077f6b1153ed3cba361401c6783bbad6114804a000cc22eb71c13e190/numpy-2.4.6-pp311-pypy311_pp73-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8ad03c0965fb3c692200e74d458ca28c1dbb4ce96f9a479a8aa041ad5fabca02", size = 16747286, upload-time = "2026-05-18T23:37:06.327Z" }, + { url = "https://files.pythonhosted.org/packages/15/ce/e5ec180bc41812edcd8daeb8639d205622c0e8c02259d8ab25a0201b3c2a/numpy-2.4.6-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:2803abfebfc990042cd494d8ce2d5f82e9d847af6d35ec486923aa19dbad5e73", size = 12504263, upload-time = "2026-05-18T23:37:09.715Z" }, +] + +[[package]] +name = "numpy" +version = "2.5.2" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.12' and sys_platform != 'darwin'", + "python_full_version >= '3.12' and sys_platform == 'darwin'", +] +sdist = { url = "https://files.pythonhosted.org/packages/9a/80/db0b4559e57ec36362bedbb05530a87fafbcb6067708c946967a41d449e7/numpy-2.5.2.tar.gz", hash = "sha256:d482d171c406ae88c5b19cad3b6a1c4c5209f886ab74bc44c2c865c23f52d860", size = 20773161, upload-time = "2026-08-09T13:48:27.962Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/69/72/dccb0aaf40972777283303919f613964227266d0c13adebb79ac124f1c3e/numpy-2.5.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:14e373cfc6387177e8409dac3c7159be8eb05cd77096cd7c950268b86f62831c", size = 16891693, upload-time = "2026-08-09T13:44:51.702Z" }, + { url = "https://files.pythonhosted.org/packages/60/2e/b5aee50a1f74ac815cf8331812cb8251e29024025de462e0c047641c614c/numpy-2.5.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4bbd96c833ecc8cc069ce518078fc8c60cb9cbfb0fea5b7a803ad65035596d03", size = 11903109, upload-time = "2026-08-09T13:44:55.501Z" }, + { url = "https://files.pythonhosted.org/packages/f3/f4/29e78102a80601cf034d4e9767022cffeca2c3b4c926e1754572ca95593d/numpy-2.5.2-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:6e8172ddfcf5cf74b811d372b570b83c60bd2de87a6fbfbebdadb4a9bd9c6cbb", size = 5350202, upload-time = "2026-08-09T13:44:58.401Z" }, + { url = "https://files.pythonhosted.org/packages/11/4b/dcd3b7eadaf4035d2c7a4289d232523a6964f602598ef7674e4bd7291f93/numpy-2.5.2-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:65f188481f1669e26f62b701e8205d19e460fa4a9b52a1414ba382330e4a3414", size = 6687736, upload-time = "2026-08-09T13:45:00.813Z" }, + { url = "https://files.pythonhosted.org/packages/e5/21/4947e0e9d6c9fc2e2ff15b8949049ee44f63adb9cacc729ab8793f97e712/numpy-2.5.2-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8ee9c4eeb8454b3660a8b53493563c3e121c2fc94fbd72b848ef814ed7b676a9", size = 15612696, upload-time = "2026-08-09T13:45:04.151Z" }, + { url = "https://files.pythonhosted.org/packages/3a/5f/62d28cf019460c7f1394105b4d49d9911a9c444cb77ab0bd95a204c5a6de/numpy-2.5.2-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3cdec01fa790a186d430433fdd4d4ffb70eed6f0eeb4bf05c8dbe2dce0a9bcb8", size = 16722264, upload-time = "2026-08-09T13:45:07.714Z" }, + { url = "https://files.pythonhosted.org/packages/14/25/3f0be4c1b9fdf5dd5e708a6806978564d7c46a055c000496309ff2a2f8af/numpy-2.5.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:7999d4ddb0c4025018373fd787510d46e04c769467af22869707b3c1cfd459ab", size = 16974396, upload-time = "2026-08-09T13:45:11.316Z" }, + { url = "https://files.pythonhosted.org/packages/22/72/6262cbdeeb45da9d971e40715f579d791603ba8ec0b5e2db1ac55454421d/numpy-2.5.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:c1f017dc0875c9209d219f97feceb7d54c2661bb243deb4114478e1295808af7", size = 18476044, upload-time = "2026-08-09T13:45:14.869Z" }, + { url = "https://files.pythonhosted.org/packages/36/33/29208b8b075bde62d26a81d14b358c42b0f69b6cabd98d4ff97f37f22b05/numpy-2.5.2-cp312-cp312-win32.whl", hash = "sha256:d6a48072864e3324e194a8fbb3c657bcc5b5c869dbc64c9537b1d5c862572c0a", size = 6072817, upload-time = "2026-08-09T13:45:17.867Z" }, + { url = "https://files.pythonhosted.org/packages/7f/b9/87fea2769fe1c47c1b5b01d8310772c9d1a85d485de7cf386ef7a3332b02/numpy-2.5.2-cp312-cp312-win_amd64.whl", hash = "sha256:28ac63476ec7651484215ee7fa15a1f78b57c14621f01e392afe17b9a1390ce4", size = 12464674, upload-time = "2026-08-09T13:45:20.734Z" }, + { url = "https://files.pythonhosted.org/packages/14/52/032b97e00461ab0809bbe4c588b035620e5a14b8cdee47ecddefc7b17d33/numpy-2.5.2-cp312-cp312-win_arm64.whl", hash = "sha256:27650bb0e7140fa3d37b9923b4803645e0b125d190f326eecfd3f4dad8e8ade1", size = 10397131, upload-time = "2026-08-09T13:45:23.73Z" }, + { url = "https://files.pythonhosted.org/packages/f5/d2/6b24738a0ef4557d189b150046cd07823c50e4273e8aebd651222e24306f/numpy-2.5.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:8e4cb9a754c8a0c62eaa88273a5fba3391f4a610d1dee893c0755da31c083f15", size = 16886595, upload-time = "2026-08-09T13:45:27.323Z" }, + { url = "https://files.pythonhosted.org/packages/65/60/f2d208d366f263f39c6e69ed309290717aab41078b6d04c9be2a84fa2a07/numpy-2.5.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:52c808f96484f5571a5cc863775ce50247c17dfb3b0361f8ed6b4b0456f80080", size = 11896845, upload-time = "2026-08-09T13:45:31.638Z" }, + { url = "https://files.pythonhosted.org/packages/3c/79/81e0bf24f4d020a2b1d5cd297a9f60c3f24eeb116f9bba5870443f7b6a4a/numpy-2.5.2-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:29d81e97f668489cba8ebfd796b9bdd453525d35dd9e162e2daec94bf3fc7740", size = 5343880, upload-time = "2026-08-09T13:45:34.373Z" }, + { url = "https://files.pythonhosted.org/packages/ba/cc/e3141cf06d1a8a2c7e107543fe1269c1d1af760d4d683c0794a4ee1127c2/numpy-2.5.2-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:afb3f0632d6b2e3ba04dbce8d1e48d321b369138b73830b5ca371a0e8d479d56", size = 6682264, upload-time = "2026-08-09T13:45:36.7Z" }, + { url = "https://files.pythonhosted.org/packages/29/f1/2a64a307d92c5d98f5255a4014eb43bb6103ee477087b61ecae44a3aa9b9/numpy-2.5.2-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0aadf13b60048d501e05fa699efaf7734e2494f3498a4c2a5521d822640324f3", size = 15609566, upload-time = "2026-08-09T13:45:39.518Z" }, + { url = "https://files.pythonhosted.org/packages/7b/44/59a1eb68e773c4098d107ef34a0dbdeca501d72ffcfbff9a7707343921ce/numpy-2.5.2-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:29b86ff8a6cc556b47ec6b64b194815cc80e6bf5eedcc6cddfd65318cb0b4eee", size = 16709995, upload-time = "2026-08-09T13:45:43.661Z" }, + { url = "https://files.pythonhosted.org/packages/8a/4c/3e54d4ddbc359a1295f8b633e8106bcd4d7d4a206e82df051bdfb3058755/numpy-2.5.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:6950c4b7dd562453090548ba7f5da7e59f57f85663f15d5dcc60e249192f7e59", size = 16972511, upload-time = "2026-08-09T13:45:47.094Z" }, + { url = "https://files.pythonhosted.org/packages/f2/9f/02e371638ebf19b66d46231e4be52999e87f32d1961b113bc45656608b22/numpy-2.5.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:b9727f472d2f3888053b8a75ab0cb94745a9de224bb5846dbadc0092101bc71d", size = 18465609, upload-time = "2026-08-09T13:45:50.808Z" }, + { url = "https://files.pythonhosted.org/packages/eb/ae/ad6645abc7a3510fe48e8ea1ab4598166f500057ef4ebf38bfad4f1577de/numpy-2.5.2-cp313-cp313-win32.whl", hash = "sha256:4f9744f9fbdcea0bc552e8f19e1f141f811a3f9bc2be2cc6e86d982cab23e3f4", size = 6070204, upload-time = "2026-08-09T13:45:54.111Z" }, + { url = "https://files.pythonhosted.org/packages/15/20/f3489f86d81ea460b2bcdceaed094142ca6579f6be0ec527b781d39afe68/numpy-2.5.2-cp313-cp313-win_amd64.whl", hash = "sha256:85aaccb24182c25df891ad0ec333585967e115269d5f1b17f2c9ae005bc96657", size = 12460532, upload-time = "2026-08-09T13:45:57.167Z" }, + { url = "https://files.pythonhosted.org/packages/d5/21/35b31dde1b283b79de828b80f876afd8c94e28fe1e9c375f89e261cc4c0d/numpy-2.5.2-cp313-cp313-win_arm64.whl", hash = "sha256:bd68ece1553d2023c09a4226d9e41c586ad2d20594d1a456186c33513d2cb3f2", size = 10396725, upload-time = "2026-08-09T13:46:00.478Z" }, + { url = "https://files.pythonhosted.org/packages/ac/f8/c3b222bf075b50afd8e949a07a15c4b312a4a84bd8102a332bcd953cbbb4/numpy-2.5.2-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:d787cf769c3baeb5f6235e778edb52c08dfa923789b5958f28e6450f96107cb1", size = 16885180, upload-time = "2026-08-09T13:46:03.939Z" }, + { url = "https://files.pythonhosted.org/packages/17/e1/2c1d4b1987795a92b5bbf7c24fe249ab96aa2573ab0d7604802c189d7b86/numpy-2.5.2-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:24b9dc2e3d84aa58523798805194e23e736f3f6ce2d1a5b92583ae734e6dbda8", size = 11907878, upload-time = "2026-08-09T13:46:07.045Z" }, + { url = "https://files.pythonhosted.org/packages/b9/ee/d08226fc858044355983a6e5b94f08ff6f3969e0a2b160a4a89f0ddb3445/numpy-2.5.2-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:9e9413326d726c2545bfa65d2c0876871e8d8386e77f992c1d426e180bbd4323", size = 5354922, upload-time = "2026-08-09T13:46:10.04Z" }, + { url = "https://files.pythonhosted.org/packages/94/f0/6d3d933056440ebbc5e6bad92065fc6c26a48a84a36b1208580e94eea76c/numpy-2.5.2-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:60e902ac295855348a5ca2ea4c89108989a9f5fddfad3dfc0a8f36b10358567e", size = 6679168, upload-time = "2026-08-09T13:46:12.275Z" }, + { url = "https://files.pythonhosted.org/packages/c4/3b/ecd49dd90033cceb2704d88ca905d4d7d89b0e8c739608754ffd325fa820/numpy-2.5.2-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:50e500dc868e9313530ce12ba470fe50ff3afe3d62993ed6eff652dacd555b65", size = 15624501, upload-time = "2026-08-09T13:46:15.322Z" }, + { url = "https://files.pythonhosted.org/packages/c7/99/461bd36dbdfac6c1c53efa370bd55a83227542d0d118f1677dbf1a3dacd5/numpy-2.5.2-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:318b9a4c845dbea06708a29c84ee429cc3065048db34cdb799047643492050ee", size = 16713701, upload-time = "2026-08-09T13:46:18.949Z" }, + { url = "https://files.pythonhosted.org/packages/f9/9c/2b251df9e8a5d647b62b0cbc1b90a91850c1cf4859ecb532fd0b4eacff6c/numpy-2.5.2-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:34c319e2963be042673fb46570501b2f06c41924e17e3563d58646b4380dfb68", size = 16986065, upload-time = "2026-08-09T13:46:23.006Z" }, + { url = "https://files.pythonhosted.org/packages/8f/25/20de43f53ff1390534a124475055a19f01fe10c920a0fd11b8e18d6d6052/numpy-2.5.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:f06571a052127dc1b4e8b83029b4d1b20daa2b64a31cdd181fc6bc774e9000eb", size = 18470031, upload-time = "2026-08-09T13:46:27.102Z" }, + { url = "https://files.pythonhosted.org/packages/56/5e/0c577ca308d6da5eb79b546ba10bbe5b60148192194e2da060913b1de4f1/numpy-2.5.2-cp314-cp314-win32.whl", hash = "sha256:2cc779226e476d1e1f08c74068c419e60f41a9e0e069c92f6671d31d5c985e98", size = 6121028, upload-time = "2026-08-09T13:46:30.046Z" }, + { url = "https://files.pythonhosted.org/packages/15/5c/7bcbd5b11f94199073320410cddcbb80cee62415bfeb540874b265c2d922/numpy-2.5.2-cp314-cp314-win_amd64.whl", hash = "sha256:7587f53dfbd5edc0f7b87c6217b4c6d2d1f2ef9c3da70bc1315e7db5f8d7ec9d", size = 12597627, upload-time = "2026-08-09T13:46:32.886Z" }, + { url = "https://files.pythonhosted.org/packages/87/bc/4d0b06fba0da90ccc75af62823cb9dcedb6c9ea0cffa058cb2c9ee773a77/numpy-2.5.2-cp314-cp314-win_arm64.whl", hash = "sha256:3e4c367352d3747784248a227fbec218e193b56f7e6692e3b64fc805478ecfdf", size = 10680414, upload-time = "2026-08-09T13:46:36.036Z" }, + { url = "https://files.pythonhosted.org/packages/cd/17/f429aac9dc08833a0d0f188eba38c532a751b1a1f2ca6018a37b455cb321/numpy-2.5.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:b879fb674276e331513fb136b78dbc6bd3c848309e0d841cfd63be3896c4cfc1", size = 12026967, upload-time = "2026-08-09T13:46:39.084Z" }, + { url = "https://files.pythonhosted.org/packages/ca/9f/d0849de96a2a4ceaa16662f18ee13eaa9c0aa418269fdc8c4857c56b11da/numpy-2.5.2-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:fd0d703772bba096843785bd38371e31bb4a0c1151497ad5739d182114a73f7f", size = 5473874, upload-time = "2026-08-09T13:46:42.075Z" }, + { url = "https://files.pythonhosted.org/packages/89/3c/8df216d4a4a5422a3de045301cf7df8ea47286d76f5cb7160b0128ac26b7/numpy-2.5.2-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:3a2f061cebd9e3d23bdcfaaded5e2293a4c6a5b60fa42df85d410a725ce621bf", size = 6789276, upload-time = "2026-08-09T13:46:44.387Z" }, + { url = "https://files.pythonhosted.org/packages/e6/3a/20d7e9891c4ddfadd6ff8d95bf4b29f353d8e1770553de2099880551dfb9/numpy-2.5.2-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6df895598c0edcb41030126c89e0f353b07d93238116143b7405e937359736c4", size = 15659154, upload-time = "2026-08-09T13:46:47.538Z" }, + { url = "https://files.pythonhosted.org/packages/aa/d6/f3aa3d2688bf501b858835c6bd087ae9b51a56ae6fca8e2b0990abd177af/numpy-2.5.2-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1ab3d4a901f844ea836c3e80bf463c6a27d7f3c14e8e292fcf28d348b25b9bce", size = 16748909, upload-time = "2026-08-09T13:46:51.442Z" }, + { url = "https://files.pythonhosted.org/packages/7d/8f/1c5cae8d2baf86ab802ae97a00be55bc7e21ebc11b12bbc33376c5f05342/numpy-2.5.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:cebc2d6dbb605a7703d59751dea4bd6b0ab127a5a4338a6f432df1936fef8b26", size = 17027685, upload-time = "2026-08-09T13:46:55.095Z" }, + { url = "https://files.pythonhosted.org/packages/5c/27/71d3467404aedc1c24ce79610f91b52b0b0f466c43a701aa56fc75c145ab/numpy-2.5.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:eaca7ff36f0f52e2111ec71f169d8fd3e889e7ddc0d2592e0d703fd8d3ce8fac", size = 18501181, upload-time = "2026-08-09T13:46:59.09Z" }, + { url = "https://files.pythonhosted.org/packages/14/2f/42921d27c40aea7e077f4a423ae509fd9220b028cd787bafefd8ab2b3a5f/numpy-2.5.2-cp314-cp314t-win32.whl", hash = "sha256:ddf47472af2e4280d79bac82304f5e80150211f1b9e614b760061d5fdfbb6eba", size = 6271085, upload-time = "2026-08-09T13:47:01.903Z" }, + { url = "https://files.pythonhosted.org/packages/75/e6/bad5f5d56de9b1971bac959963dda276d35c40f1854475005434bbe08692/numpy-2.5.2-cp314-cp314t-win_amd64.whl", hash = "sha256:44ef9675d908e65f9953063837c3277730f3f4437615a4cdab67b366cabaf884", size = 12787971, upload-time = "2026-08-09T13:47:04.963Z" }, + { url = "https://files.pythonhosted.org/packages/df/05/f608795cb34391acd67e38d94a3c36abd8d8576293a3a80727d7595c372c/numpy-2.5.2-cp314-cp314t-win_arm64.whl", hash = "sha256:eaa088384c46f519dacb93b7ec483a6d6b19a4a2085ae4f25ab9b1c43d387d1e", size = 10750306, upload-time = "2026-08-09T13:47:07.976Z" }, + { url = "https://files.pythonhosted.org/packages/33/c6/28de0191c5f82b7d42a0a51390ba98587048aa93a39fafb05bdbe6e8d00c/numpy-2.5.2-cp315-cp315-macosx_10_15_x86_64.whl", hash = "sha256:078f9b027b478c9379b9677babbf0f8b8f1ecfada27636d7b9a93990c638739f", size = 16885274, upload-time = "2026-08-09T13:47:11.439Z" }, + { url = "https://files.pythonhosted.org/packages/dd/d1/973ca116000d244897e468ea1aff30b589e5022e3c8744b71706fe33bd57/numpy-2.5.2-cp315-cp315-macosx_11_0_arm64.whl", hash = "sha256:50a68f4bacd8a2b33d8da3d2269d0d78500f86ea582e4786dc10f5ef2c2c6842", size = 11907846, upload-time = "2026-08-09T13:47:15.128Z" }, + { url = "https://files.pythonhosted.org/packages/78/d9/8c4b3937ef204cb2fd88d389ccd0f265a2ffb11f35a01d2064cf46714bd6/numpy-2.5.2-cp315-cp315-macosx_14_0_arm64.whl", hash = "sha256:e79aba74ffaf5f78a050d777c184cddf8fdffabab38acf5f3ef1fecbc17895d6", size = 5354892, upload-time = "2026-08-09T13:47:18.07Z" }, + { url = "https://files.pythonhosted.org/packages/74/9b/b6ee65ea2999fdb7023935e108e6fb776ee4082aa15f159acfa857e578c8/numpy-2.5.2-cp315-cp315-macosx_14_0_x86_64.whl", hash = "sha256:9a0731745a72a184490a582fb4af2533512bd071ace67785b5fdffc0ae58dce8", size = 6679309, upload-time = "2026-08-09T13:47:20.456Z" }, + { url = "https://files.pythonhosted.org/packages/43/f3/acb18d8b137a393c8e7803a8c994c9e64bde3930692a69d826993113a159/numpy-2.5.2-cp315-cp315-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4ec954036759bcee3aa484f8603bd9c14f3e776293b85578b8734c2d72777c69", size = 15625850, upload-time = "2026-08-09T13:47:24.365Z" }, + { url = "https://files.pythonhosted.org/packages/a9/bf/a8e9bb0db815a0e265b5744ebedd3af0bd5faad8604e5b50a1cd012f3c91/numpy-2.5.2-cp315-cp315-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:dc649493697006bc90614a5f0bbc8cb3cb1866715c474e473694968d7e6b99ab", size = 16713664, upload-time = "2026-08-09T13:47:27.965Z" }, + { url = "https://files.pythonhosted.org/packages/0c/c3/6e913736b3dd6582344af32418b5fb9dab34282e8a8174ae1d54ceb0fc13/numpy-2.5.2-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:cf7de32f486e4ac9e2d93b810f9e9ac72a728dd46a32a0bb403222f27f653514", size = 16986749, upload-time = "2026-08-09T13:47:31.541Z" }, + { url = "https://files.pythonhosted.org/packages/80/09/7d3b23eff5c7428ef6c01e6f7052bb60d504c4d33e317b36b8959c24ad97/numpy-2.5.2-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:2ffa7bacab3e2ee1b19ed31766bb60bb380b68c23f051e199c5cc598afd68710", size = 18470495, upload-time = "2026-08-09T13:47:35.364Z" }, + { url = "https://files.pythonhosted.org/packages/a5/a4/68a321d825374f6eb677ffe8ef8c6b9a328304e6fd2e39d9530822776607/numpy-2.5.2-cp315-cp315-win32.whl", hash = "sha256:6b588cc8f902d6bff201c19fd00c43ab8545671e3554d014e12e14139e5e8617", size = 6120696, upload-time = "2026-08-09T13:47:38.561Z" }, + { url = "https://files.pythonhosted.org/packages/c8/23/deafbb1700f79fae9cd1e91220f133d124cc267de1b584da3fbf6db2f6cd/numpy-2.5.2-cp315-cp315-win_amd64.whl", hash = "sha256:07d4e89f3a9ab0a9ba24264ccdb642b3dd951b2281e8883a5481a4aa79cc31a7", size = 12597324, upload-time = "2026-08-09T13:47:41.401Z" }, + { url = "https://files.pythonhosted.org/packages/33/cd/3272ba105e3bbbdaeb11357eda31e7a6825ffe159e8171665660299a948f/numpy-2.5.2-cp315-cp315-win_arm64.whl", hash = "sha256:a610dc7e3c52edd39c2bc2375ff9c3fd59cb3ad00e4472d36f83bc1457145788", size = 10680466, upload-time = "2026-08-09T13:47:44.873Z" }, + { url = "https://files.pythonhosted.org/packages/0e/0e/58370637b1bb70a5c9ce2b43f4b521ccb224e36ccb76a6596b17ae4b447c/numpy-2.5.2-cp315-cp315t-macosx_10_15_x86_64.whl", hash = "sha256:40f4d451aed46a8046a1aae41c4e55fb3612273df9c502480135e1501576a34b", size = 16993947, upload-time = "2026-08-09T13:47:48.97Z" }, + { url = "https://files.pythonhosted.org/packages/10/93/2abcb807712b289d6d60fe4cf30532f98974a8396d885650f3ba5a13026e/numpy-2.5.2-cp315-cp315t-macosx_11_0_arm64.whl", hash = "sha256:c081cbe16ba1ab53078e5ff29013621e33c509eedab055775d956427712c236e", size = 12025331, upload-time = "2026-08-09T13:47:52.646Z" }, + { url = "https://files.pythonhosted.org/packages/8b/3a/2898e003a5fbaf87e76c039b4ee1f5eb390471b4ffe74887c1f34c4e791e/numpy-2.5.2-cp315-cp315t-macosx_14_0_arm64.whl", hash = "sha256:0090ccdd57ec2703e9b49d0bf554767370581c1dd0a6b2bb2b2d9def317d042a", size = 5472336, upload-time = "2026-08-09T13:47:55.403Z" }, + { url = "https://files.pythonhosted.org/packages/61/a5/23f69d07c544597b29758b31b55c27dc9d541012a2c1496189fef702aec2/numpy-2.5.2-cp315-cp315t-macosx_14_0_x86_64.whl", hash = "sha256:6a9bb119fb8dd21ba30b3f0e555b7e2b081bd9883af21ec9c1c633d161cda3a8", size = 6788387, upload-time = "2026-08-09T13:47:58.192Z" }, + { url = "https://files.pythonhosted.org/packages/15/ea/c0dbdbcf22f43782510a3e492dd3da73c6112b69cac8929d16d127536fc4/numpy-2.5.2-cp315-cp315t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a839318485284a6fb31be4f8f2c91c8f2cb22f4543c4a8903f12b0671ffe07cc", size = 15667096, upload-time = "2026-08-09T13:48:01.562Z" }, + { url = "https://files.pythonhosted.org/packages/fc/5e/29c73c31748cdb0f7566642125ba17fd5b56780cddf891b085dab27e4466/numpy-2.5.2-cp315-cp315t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ba0a474801b8dc67b66bf465548abc90e82b44d2611b5770f33008dcabffe8ec", size = 16751730, upload-time = "2026-08-09T13:48:05.706Z" }, + { url = "https://files.pythonhosted.org/packages/47/95/02501e8454796bb58dadf7a99d3181e0b464bf264e1003039572f9779fac/numpy-2.5.2-cp315-cp315t-musllinux_1_2_aarch64.whl", hash = "sha256:0a4035ae1129ff8777f08bfbd44f1e5d8e9c049ce0c2dd78fc0d92c13e7251c0", size = 17038686, upload-time = "2026-08-09T13:48:09.627Z" }, + { url = "https://files.pythonhosted.org/packages/0e/b5/53a681d91b5c82687067d8ea5035e02d917b5509d6f334cb06484a954714/numpy-2.5.2-cp315-cp315t-musllinux_1_2_x86_64.whl", hash = "sha256:77843ca236b777e67f8d6b3660ea116e499612703a0ecd7093f316201eb9d8e2", size = 18507727, upload-time = "2026-08-09T13:48:13.744Z" }, + { url = "https://files.pythonhosted.org/packages/42/06/6e11443f7b64ee376c860506091103bf68f92d2cab9e8d96d4501babf07c/numpy-2.5.2-cp315-cp315t-win32.whl", hash = "sha256:7354826bc6f8f69402e9b7fe28d15fcd34feebd74f856f111585c5b0c9fb0251", size = 6269775, upload-time = "2026-08-09T13:48:17.543Z" }, + { url = "https://files.pythonhosted.org/packages/f1/18/195d6b86cd72dbbc501edfa778005fa6b87afd34c153e46028cd3a0938f4/numpy-2.5.2-cp315-cp315t-win_amd64.whl", hash = "sha256:e5651f3f87add730ee6608d915009e19c911fba0cb000c7e3ea994b7d768eb12", size = 12782559, upload-time = "2026-08-09T13:48:21.023Z" }, + { url = "https://files.pythonhosted.org/packages/b4/07/458c344f0f0c178f4481dad5cca790626ffe4c34eabf9467069d06ee4999/numpy-2.5.2-cp315-cp315t-win_arm64.whl", hash = "sha256:5f8e00be2ec6f45f4e8a41a527f68d44a7d96fee92a650e4d8b1326f77f61e6e", size = 10748103, upload-time = "2026-08-09T13:48:24.21Z" }, +] + +[[package]] +name = "nvidia-cuda-crt" +version = "13.3.73" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fa/41/2089e411507d66458d67208bdd1bc562d492bb6458c3d2aea4603072a219/nvidia_cuda_crt-13.3.73-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:60aacc0b5e1e8b40c62abe4d1ab16440add91b99bd2f17f62dd091586b73d166", size = 157353, upload-time = "2026-06-29T16:42:38.163Z" }, + { url = "https://files.pythonhosted.org/packages/7e/ce/16d76f4b5b3f7460f5ebd17516685495c149c66651ffdd381f90e4d4e65c/nvidia_cuda_crt-13.3.73-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:df14a17ae1c5c3171265411212246654d780f89344ea85344466c6b955247543", size = 157352, upload-time = "2026-06-29T16:43:09.209Z" }, + { url = "https://files.pythonhosted.org/packages/49/b3/6791ffba6f4b8e0d3ed875285aad8078ee407afa464ecd934ae298c205b1/nvidia_cuda_crt-13.3.73-py3-none-win_amd64.whl", hash = "sha256:af04e75148db1f0eea30958f33a9ec5a5a2dc2afa99ca4323f9a93b840602ca5", size = 158286, upload-time = "2026-06-29T17:09:28.621Z" }, +] + +[[package]] +name = "nvidia-cuda-nvcc" +version = "13.3.73" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cuda-crt" }, + { name = "nvidia-cuda-runtime" }, + { name = "nvidia-nvvm" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/5c/14/9f5cdc994d5431e2f08f62ffe34509e7feabd1f2e18517e2d7720c6ff0fd/nvidia_cuda_nvcc-13.3.73-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:70f250825355d2c3aa6c7a972a0ec00f020bad66d2679e527eb4336301c904aa", size = 39515578, upload-time = "2026-06-29T16:47:40.318Z" }, + { url = "https://files.pythonhosted.org/packages/83/19/e46ef3597ba47a9f8a91ab24533db42a600b659fc418dbe4af0b630bcb41/nvidia_cuda_nvcc-13.3.73-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f483af83166c4fa356a21606076d553b0b4ceaebbd9912e537545080db695bdd", size = 44942138, upload-time = "2026-06-29T16:48:13.615Z" }, + { url = "https://files.pythonhosted.org/packages/79/89/97eb797bb8bdee1d4e74069d072c24b79ae90c012fa3b539f2a7ccecf6cf/nvidia_cuda_nvcc-13.3.73-py3-none-win_amd64.whl", hash = "sha256:3d9da631bcac3dee49d1357b84cd05abe56aa3ccf76b05a7df8a80ef78addcb5", size = 32536529, upload-time = "2026-06-29T17:11:25.455Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime" +version = "13.3.29" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5f/e5/c1a221c8e6fecd071b80ea44c20fc253ae24f56e15e3f77cfbc3fb76e724/nvidia_cuda_runtime-13.3.29-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:73291e19c9dd919c140c91bda2f80b0eca487da5ee30a086ef7bc4918ecb90ea", size = 2356574, upload-time = "2026-05-26T16:29:56.333Z" }, + { url = "https://files.pythonhosted.org/packages/97/be/5699b6e642b372f7d24c59c2f41383e2696825e20bab85f7399c7c6a56f7/nvidia_cuda_runtime-13.3.29-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e04420616e72f563167a7733272992d7e6df6dc5cb54b2f94f9f1520ea9e30c1", size = 2339786, upload-time = "2026-05-26T16:30:21.584Z" }, + { url = "https://files.pythonhosted.org/packages/d2/27/b53a5e0397842a5c11f0e1a39d4e5b2f22638a4126e83b3c4e196f62c969/nvidia_cuda_runtime-13.3.29-py3-none-win_amd64.whl", hash = "sha256:0667ec61c3d897388efa305ed4f7609ace88849a753ba9c6311d06dca55fff4f", size = 2630354, upload-time = "2026-05-26T17:00:05.389Z" }, +] + +[[package]] +name = "nvidia-cuda-tileiras" +version = "13.3.36" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cuda-nvcc" }, + { name = "nvidia-nvjitlink" }, + { name = "nvidia-nvvm" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/4b/bd/89c65b2067a2a707cc92fad6d34e96ae53a495805f563642244649f9cee2/nvidia_cuda_tileiras-13.3.36-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:91510f8f9b4f4b3f3a365c43218b62873b5ae1474db6b32eb3ece7f0934496af", size = 36181023, upload-time = "2026-05-26T16:41:51.373Z" }, + { url = "https://files.pythonhosted.org/packages/80/39/c51197d0754f7a2b4bb21e0939c6940c5d4258aa0ce89974c404fc10977c/nvidia_cuda_tileiras-13.3.36-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9221618a3f6422efcb9ae38225d1db7bf0f4ca211f8fdc98a3cd6a81a8f00eca", size = 37050964, upload-time = "2026-05-26T16:42:26.833Z" }, + { url = "https://files.pythonhosted.org/packages/54/b1/9dcc1aa140205c8f9ecdb671b488e1189462d870e5ed0eac02522722cb9a/nvidia_cuda_tileiras-13.3.36-py3-none-win_amd64.whl", hash = "sha256:0dd286086c6d273826218d02c076ddea8e285b50ee223bd1429756b706b7bbc7", size = 29715011, upload-time = "2026-05-26T17:05:29.974Z" }, +] + +[[package]] +name = "nvidia-nvjitlink" +version = "13.3.33" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f0/ee/580ca6f29dcab0221db8706badca1bbbb084f1975c4d4e83329c3a7e31f0/nvidia_nvjitlink-13.3.33-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:26a6de7fb4c8fdaa7703d3dad720d6d427ddfea5c48a528fd97c11733ad830e5", size = 40742423, upload-time = "2026-05-26T16:54:51.613Z" }, + { url = "https://files.pythonhosted.org/packages/69/30/45414e35ff2eee7db3da037e5707037ccf9d2b5218ffbdb055ea4d5aa98a/nvidia_nvjitlink-13.3.33-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ce48b37dfeb3cb1eae4cf85adacb47d7a6539ea2272870c9a3628ce275c2037e", size = 39168635, upload-time = "2026-05-26T16:54:13.906Z" }, + { url = "https://files.pythonhosted.org/packages/67/f2/ec9c05a108095828dfc58840978c627b3c313fdf2a567c6de9ffbbb46901/nvidia_nvjitlink-13.3.33-py3-none-win_amd64.whl", hash = "sha256:4297ee49639b4f2e07255a1d69b3acc7ab2d011bb892b403e91ac98368962e3b", size = 37766359, upload-time = "2026-05-26T17:11:28.96Z" }, +] + +[[package]] +name = "nvidia-nvvm" +version = "13.3.73" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f3/e7/ff646aa6015c7e6d12aad234e68925c87b6681d8d18c3ac40535994a3b0d/nvidia_nvvm-13.3.73-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:0e28e0858a3475e11ac67d35301cd5bf82666a1c0dc4ec4e80ceaf3a5fd1dea8", size = 69250424, upload-time = "2026-06-29T17:08:07.453Z" }, + { url = "https://files.pythonhosted.org/packages/2f/05/35754a7105563fd9b496e5ee8e1acd986aef8258760c3cbccf419aee861a/nvidia_nvvm-13.3.73-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e2bcdd5783b5481445f1f0e7170cb836cc0d72999839ba850bbba6dc97b76bb8", size = 66984478, upload-time = "2026-06-29T17:07:43.765Z" }, + { url = "https://files.pythonhosted.org/packages/ad/6b/d5756f485012b920475cbc01457c1b9a7d0485bfb04b92598c5e1ef3e9ab/nvidia_nvvm-13.3.73-py3-none-win_amd64.whl", hash = "sha256:b5c91dfa59ee4cee90b2dfb19c6203f31c914b9c9b5ca10726c2da7cf8ed401d", size = 59981103, upload-time = "2026-06-29T17:21:43.334Z" }, +] + +[[package]] +name = "packaging" +version = "26.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7d/fa/3944b40b07da9ce895c0e6303a5ab7d53da063554f534556b134a54d6093/packaging-26.3.tar.gz", hash = "sha256:94edc256424af38762eb31306eed28beb9f0efc50a8837492c9d6fd6004aed79", size = 313412, upload-time = "2026-08-04T18:15:28.737Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/34/ba1c580383c9eada3711951fef0795c80b829a078d72188184bcab9dd527/packaging-26.3-py3-none-any.whl", hash = "sha256:d7193f7c8e4e93f444fde0262bf90af30e16fa0ad0ad44cb553c87339b23cd1c", size = 129956, upload-time = "2026-08-04T18:15:27.159Z" }, +] + +[[package]] +name = "pybind11" +version = "3.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/76/f3/95b0f40b31df41dbfe6bb0857419c9442c15839cbac4796f1c26ae0b6081/pybind11-3.1.0.tar.gz", hash = "sha256:a1cc06b524ab3edca51f8ad3895f9c4fa20b8b19283173dff4ae781449dc9639", size = 603746, upload-time = "2026-08-06T23:33:00.675Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/33/fd/8762f7ee3e4e4be6d1d846cffb4916dd9bb02b2f800d3603718a0efe494c/pybind11-3.1.0-py3-none-any.whl", hash = "sha256:b8488090f8acffbcb6b5d6a85571a6827a0a2981ffb75e5a0b27b87c4a6b7dd0", size = 319402, upload-time = "2026-08-06T23:32:59.047Z" }, +] + +[[package]] +name = "setuptools" +version = "84.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6d/44/f5da03a8ef95d369145c5bb53050e7877c9f3d312e128605fd9504829143/setuptools-84.0.0.tar.gz", hash = "sha256:f4695c21257f0d9b537ec2692c941d02ee143b7cc1276941349a546573b2ef73", size = 1168449, upload-time = "2026-08-08T18:27:58.365Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/95/9c/c510029fc6ef33a6275cd2c5d3cecd6613dfd6aa401d57c54f1c18852ccf/setuptools-84.0.0-py3-none-any.whl", hash = "sha256:51a52592b3b99e102b609654876bd65f19f999935166d1352678931132b0c670", size = 818216, upload-time = "2026-08-08T18:27:56.719Z" }, +] + +[[package]] +name = "sympy" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mpmath" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, +] + +[[package]] +name = "torch" +version = "2.13.0" +source = { registry = "https://download.pytorch.org/whl/cpu" } +resolution-markers = [ + "python_full_version >= '3.12' and sys_platform == 'darwin'", + "python_full_version < '3.12' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "filelock", marker = "sys_platform == 'darwin'" }, + { name = "fsspec", marker = "sys_platform == 'darwin'" }, + { name = "jinja2", marker = "sys_platform == 'darwin'" }, + { name = "networkx", marker = "sys_platform == 'darwin'" }, + { name = "setuptools", marker = "sys_platform == 'darwin'" }, + { name = "sympy", marker = "sys_platform == 'darwin'" }, + { name = "typing-extensions", marker = "sys_platform == 'darwin'" }, +] +wheels = [ + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:e76f9bcecc52b8ff711239a2f7547d5353df95878ab232f0773c1d95928b92f8", upload-time = "2026-07-08T12:26:13Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:2fe228aba290d14b9f31b049be550dbd469c3fd3013d7a19705b30454da97027", upload-time = "2026-07-08T12:26:18Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:33449899ce5496c1b84b4853179d94fd102028ae1407314d9fb956bb79e70d09", upload-time = "2026-07-08T12:26:23Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:d849b390e07d8d333ce8ecaf91b273c656c598379a19c9acf1318a883f6b391c", upload-time = "2026-07-08T12:26:28Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:c28def70706c2f9ecc752574766e8ae4da9b810ab6676b611166761a78a9f1e1", upload-time = "2026-07-08T12:26:33Z" }, +] + +[[package]] +name = "torch" +version = "2.13.0+cpu" +source = { registry = "https://download.pytorch.org/whl/cpu" } +resolution-markers = [ + "python_full_version >= '3.12' and sys_platform != 'darwin'", + "python_full_version < '3.12' and sys_platform != 'darwin'", +] +dependencies = [ + { name = "filelock", marker = "sys_platform != 'darwin'" }, + { name = "fsspec", marker = "sys_platform != 'darwin'" }, + { name = "jinja2", marker = "sys_platform != 'darwin'" }, + { name = "networkx", marker = "sys_platform != 'darwin'" }, + { name = "setuptools", marker = "sys_platform != 'darwin'" }, + { name = "sympy", marker = "sys_platform != 'darwin'" }, + { name = "typing-extensions", marker = "sys_platform != 'darwin'" }, +] +wheels = [ + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp311-cp311-linux_s390x.whl", hash = "sha256:6e9817dbdf5ea76789babd46e457eac5bf14ff566cf85f8addbfdff2d56601ce", upload-time = "2026-07-08T19:27:52Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:84453b69508ec79902f899c5ed9495acb9e2bbe9fda5f1d5d6f19e3c3842e1a7", upload-time = "2026-07-08T19:28:03Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:6746dbcbeb526eb61330b76b41ff1b4eb848951103a892eeb080dfa2b264667b", upload-time = "2026-07-08T19:28:16Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp311-cp311-win_amd64.whl", hash = "sha256:10717d8b3b67c45a4788bf7ffc0bab1ea1e5ebbedd24466be6100102d141fac1", upload-time = "2026-07-08T19:28:27Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp311-cp311-win_arm64.whl", hash = "sha256:2b3d093abd919ad934c43d47e73ba63ceba7cbd7269fc2e9c1e4fc29e8fe45fa", upload-time = "2026-07-08T19:28:33Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp312-cp312-linux_s390x.whl", hash = "sha256:ffadde149901c8afa138daa38d898264003cfcf1a3336ca5cd964b5af227d867", upload-time = "2026-07-08T19:28:41Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:6f307c2c32d764ffc6ff6893b801fad6d4752f3e67966cb8abf1843427c02604", upload-time = "2026-07-08T19:28:51Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:4ca4a9394b0c771238a4f73590fdbbc4debad85ed0fa63d026ae1b085da7d6e2", upload-time = "2026-07-08T19:29:03Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp312-cp312-win_amd64.whl", hash = "sha256:a8b450c1e58e5800e5b4691dac412f8d2d65a1dc3298166f91596603a3531e6f", upload-time = "2026-07-08T19:29:15Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp312-cp312-win_arm64.whl", hash = "sha256:fa0762705b933624d59f6823db9ce7ec2e35b3e1e9c319c9db51fbeecfc3e319", upload-time = "2026-07-08T19:29:21Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp313-cp313-linux_s390x.whl", hash = "sha256:966d020354f465672dc7dd10d3a5c6cd17d7eb48620aa1d265b48a1f78f06898", upload-time = "2026-07-08T19:29:30Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:0b8f7d0423027ae8b90c7977c627f3379f325363a08224dffad9b4b2d684a83d", upload-time = "2026-07-08T19:29:40Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:3fbf9c9d1f3c10c2d59d04aca426dee9ccc6ceb32d255c61e93acc3b4f75fae6", upload-time = "2026-07-08T19:29:54Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp313-cp313-win_amd64.whl", hash = "sha256:a17ff48608634db245e17e8bb00a9558554a49aeb1e4f5fe6cd039af2a10515b", upload-time = "2026-07-08T19:30:05Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp313-cp313-win_arm64.whl", hash = "sha256:ac7aaf322be4777765a53bed7264a214dd81b3a1d276b93150515a3c5f75e4b0", upload-time = "2026-07-08T19:30:12Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp314-cp314-linux_s390x.whl", hash = "sha256:dec241fef3984c0d1edadd1f58708e218d4eae881ceef7bc10cf9964d41b68b9", upload-time = "2026-07-08T19:30:20Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:ca021f9eb2f8345c83fa03e3a04587308afb8df71bd472670b3ece00df58621c", upload-time = "2026-07-08T19:30:32Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:d20fa53ee744502fa4c69818a720b05ca0d37abd055d4f6e66cae155114bc691", upload-time = "2026-07-08T19:30:45Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp314-cp314-win_amd64.whl", hash = "sha256:e2e5134decf00e218da62318f3dc5df156231d367871918e91eba95ab0ad43ab", upload-time = "2026-07-08T19:30:58Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp314-cp314t-linux_s390x.whl", hash = "sha256:991cc14b39e751122c01f017be6448533989868731cb5eecd1006893d26787c2", upload-time = "2026-07-08T19:31:09Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:7b8d26e29bceafbdaa8d63bfe7612f23875b5af2cc07e13f809c3ed890bbe1d8", upload-time = "2026-07-08T19:31:21Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:b222c15a0fc2ce207d1c1a59700b46c8fa6748df1f447ad11e5c870dde0933d9", upload-time = "2026-07-08T19:31:35Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp314-cp314t-win_amd64.whl", hash = "sha256:a43376bd094124ef626bfdd3d4c2c62eacb0b5ddc99776f4a32d4fd16f1f3420", upload-time = "2026-07-08T19:31:48Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp315-cp315-manylinux_2_28_aarch64.whl", hash = "sha256:8eb5002ca81af00ae69b57540f615b58b8ae922b6d4848176b366a52bd2196e6", upload-time = "2026-07-08T19:32:00Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp315-cp315-manylinux_2_28_x86_64.whl", hash = "sha256:1a3a35229fdc13446b4eab50e7fcf9399ff941e89a3b761497786297a5d8dde5", upload-time = "2026-07-08T19:32:16Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp315-cp315t-manylinux_2_28_aarch64.whl", hash = "sha256:8e109528e6bab044815daebaf71770fbaace3a66ef1c816cb55c875350f78a60", upload-time = "2026-07-08T19:32:30Z" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.13.0%2Bcpu-cp315-cp315t-manylinux_2_28_x86_64.whl", hash = "sha256:222a6681467cc7f6f05cd3068dfbc603def3a1e46d1d4620c1c8cdf6178bd563", upload-time = "2026-07-08T19:32:44Z" }, +] [[package]] name = "typing-extensions" -version = "4.15.0" +version = "4.16.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f6/cc/6253133b5bb138fc3306cebfbda2c520f545d36b5be2c7255cc528bb45d6/typing_extensions-4.16.0.tar.gz", hash = "sha256:dc983d19a509c94dba722ee6abd33940f7c05a89e243c47e907eb4db6f1a43e5", size = 113555, upload-time = "2026-07-02T08:40:05.92Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/d3/b8441a820a491ddfc024b0b0cf0393375b75ea13866d9c66727e54c2fc80/typing_extensions-4.16.0-py3-none-any.whl", hash = "sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8", size = 45571, upload-time = "2026-07-02T08:40:04.659Z" }, +] + +[[package]] +name = "wheel" +version = "0.48.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "packaging" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d0/20/50ed6bdf27dec98b568a8ae25dc599f35baa3d9709f9e83fd1edb56b9a90/wheel-0.48.0.tar.gz", hash = "sha256:94800765601e9171bf5d58d066e640662842bcedcbab982b2c90787a2c987322", size = 66471, upload-time = "2026-08-11T22:02:27.327Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, + { url = "https://files.pythonhosted.org/packages/2e/29/69cfbb602cd91690c55d38ba9fe53e6a7e76a6fa647bf38f19c138d25449/wheel-0.48.0-py3-none-any.whl", hash = "sha256:3217dcc807155e45db462d7ef2431f5ddda0d7273b700d05a67b271ceb1287ab", size = 33320, upload-time = "2026-08-11T22:02:26.1Z" }, ]